WompUniversity/Inpaint-Anything-no-errors

1import sys2import argparse3import numpy as np4from pathlib import Path5from matplotlib import pyplot as plt6from typing import Any, Dict, List7import torch8 9sys.path.insert(0, str(Path(__file__).resolve().parent / "third_party" / "segment-anything"))10 11from segment_anything import SamPredictor, sam_model_registry12from utils import load_img_to_array, save_array_to_img, dilate_mask, \13    show_mask, show_points14 15 16def predict_masks_with_sam(17        img: np.ndarray,18        point_coords: List[List[float]],19        point_labels: List[int],20        model_type: str,21        ckpt_p: str,22        device="cuda"23):24    point_coords = np.array(point_coords)25    point_labels = np.array(point_labels)26    sam = sam_model_registry[model_type](checkpoint=ckpt_p)27    sam.to(device=device)28    predictor = SamPredictor(sam)29 30    predictor.set_image(img)31    masks, scores, logits = predictor.predict(32        point_coords=point_coords,33        point_labels=point_labels,34        multimask_output=True,35    )36    return masks, scores, logits37 38 39def setup_args(parser):40    parser.add_argument(41        "--input_img", type=str, required=True,42        help="Path to a single input img",43    )44    parser.add_argument(45        "--point_coords", type=float, nargs='+', required=True,46        help="The coordinate of the point prompt, [coord_W coord_H].",47    )48    parser.add_argument(49        "--point_labels", type=int, nargs='+', required=True,50        help="The labels of the point prompt, 1 or 0.",51    )52    parser.add_argument(53        "--dilate_kernel_size", type=int, default=None,54        help="Dilate kernel size. Default: None",55    )56    parser.add_argument(57        "--output_dir", type=str, required=True,58        help="Output path to the directory with results.",59    )60    parser.add_argument(61        "--sam_model_type", type=str,62        default="vit_h", choices=['vit_h', 'vit_l', 'vit_b'],63        help="The type of sam model to load. Default: 'vit_h"64    )65    parser.add_argument(66        "--sam_ckpt", type=str, required=True,67        help="The path to the SAM checkpoint to use for mask generation.",68    )69 70 71if __name__ == "__main__":72    """Example usage:73    python sam_segment.py \74        --input_img FA_demo/FA1_dog.png \75        --point_coords 750 500 \76        --point_labels 1 \77        --dilate_kernel_size 15 \78        --output_dir ./results \79        --sam_model_type "vit_h" \80        --sam_ckpt sam_vit_h_4b8939.pth81    """82    parser = argparse.ArgumentParser()83    setup_args(parser)84    args = parser.parse_args(sys.argv[1:])85    device = "cuda" if torch.cuda.is_available() else "cpu"86 87    img = load_img_to_array(args.input_img)88 89    masks, _, _ = predict_masks_with_sam(90        img,91        [args.point_coords],92        args.point_labels,93        model_type=args.sam_model_type,94        ckpt_p=args.sam_ckpt,95        device=device,96    )97    masks = masks.astype(np.uint8) * 25598 99    # dilate mask to avoid unmasked edge effect100    if args.dilate_kernel_size is not None:101        masks = [dilate_mask(mask, args.dilate_kernel_size) for mask in masks]102 103    # visualize the segmentation results104    img_stem = Path(args.input_img).stem105    out_dir = Path(args.output_dir) / img_stem106    out_dir.mkdir(parents=True, exist_ok=True)107    for idx, mask in enumerate(masks):108        # path to the results109        mask_p = out_dir / f"mask_{idx}.png"110        img_points_p = out_dir / f"with_points.png"111        img_mask_p = out_dir / f"with_{Path(mask_p).name}"112 113        # save the mask114        save_array_to_img(mask, mask_p)115 116        # save the pointed and masked image117        dpi = plt.rcParams['figure.dpi']118        height, width = img.shape[:2]119        plt.figure(figsize=(width/dpi/0.77, height/dpi/0.77))120        plt.imshow(img)121        plt.axis('off')122        show_points(plt.gca(), [args.point_coords], args.point_labels,123                    size=(width*0.04)**2)124        plt.savefig(img_points_p, bbox_inches='tight', pad_inches=0)125        show_mask(plt.gca(), mask, random_color=False)126        plt.savefig(img_mask_p, bbox_inches='tight', pad_inches=0)127        plt.close()