# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. import cv2 # type: ignore from segment_anything import SamAutomaticMaskGenerator, sam_model_registry import argparse import json import os from typing import Any, Dict, List parser = argparse.ArgumentParser( description=( "Runs automatic mask generation on an input image or directory of images, " "and outputs masks as either PNGs or COCO-style RLEs. Requires open-cv, " "as well as pycocotools if saving in RLE format." ) ) parser.add_argument( "--input", type=str, required=True, help="Path to either a single input image or folder of images.", ) parser.add_argument( "--output", type=str, required=True, help=( "Path to the directory where masks will be output. Output will be either a folder " "of PNGs per image or a single json with COCO-style masks." ), ) parser.add_argument( "--model-type", type=str, default="default", help="The type of model to load, in ['default', 'vit_l', 'vit_b']", ) parser.add_argument( "--checkpoint", type=str, required=True, help="The path to the SAM checkpoint to use for mask generation.", ) parser.add_argument("--device", type=str, default="cuda", help="The device to run generation on.") parser.add_argument( "--convert-to-rle", action="store_true", help=( "Save masks as COCO RLEs in a single json instead of as a folder of PNGs. " "Requires pycocotools." ), ) amg_settings = parser.add_argument_group("AMG Settings") amg_settings.add_argument( "--points-per-side", type=int, default=None, help="Generate masks by sampling a grid over the image with this many points to a side.", ) amg_settings.add_argument( "--points-per-batch", type=int, default=None, help="How many input points to process simultaneously in one batch.", ) amg_settings.add_argument( "--pred-iou-thresh", type=float, default=None, help="Exclude masks with a predicted score from the model that is lower than this threshold.", ) amg_settings.add_argument( "--stability-score-thresh", type=float, default=None, help="Exclude masks with a stability score lower than this threshold.", ) amg_settings.add_argument( "--stability-score-offset", type=float, default=None, help="Larger values perturb the mask more when measuring stability score.", ) amg_settings.add_argument( "--box-nms-thresh", type=float, default=None, help="The overlap threshold for excluding a duplicate mask.", ) amg_settings.add_argument( "--crop-n-layers", type=int, default=None, help=( "If >0, mask generation is run on smaller crops of the image to generate more masks. " "The value sets how many different scales to crop at." ), ) amg_settings.add_argument( "--crop-nms-thresh", type=float, default=None, help="The overlap threshold for excluding duplicate masks across different crops.", ) amg_settings.add_argument( "--crop-overlap-ratio", type=int, default=None, help="Larger numbers mean image crops will overlap more.", ) amg_settings.add_argument( "--crop-n-points-downscale-factor", type=int, default=None, help="The number of points-per-side in each layer of crop is reduced by this factor.", ) amg_settings.add_argument( "--min-mask-region-area", type=int, default=None, help=( "Disconnected mask regions or holes with area smaller than this value " "in pixels are removed by postprocessing." ), ) def write_masks_to_folder(masks: List[Dict[str, Any]], path: str) -> None: header = "id,area,bbox_x0,bbox_y0,bbox_w,bbox_h,point_input_x,point_input_y,predicted_iou,stability_score,crop_box_x0,crop_box_y0,crop_box_w,crop_box_h" # noqa metadata = [header] for i, mask_data in enumerate(masks): mask = mask_data["segmentation"] filename = f"{i}.png" cv2.imwrite(os.path.join(path, filename), mask * 255) mask_metadata = [ str(i), str(mask_data["area"]), *[str(x) for x in mask_data["bbox"]], *[str(x) for x in mask_data["point_coords"][0]], str(mask_data["predicted_iou"]), str(mask_data["stability_score"]), *[str(x) for x in mask_data["crop_box"]], ] row = ",".join(mask_metadata) metadata.append(row) metadata_path = os.path.join(path, "metadata.csv") with open(metadata_path, "w") as f: f.write("\n".join(metadata)) return def get_amg_kwargs(args): amg_kwargs = { "points_per_side": args.points_per_side, "points_per_batch": args.points_per_batch, "pred_iou_thresh": args.pred_iou_thresh, "stability_score_thresh": args.stability_score_thresh, "stability_score_offset": args.stability_score_offset, "box_nms_thresh": args.box_nms_thresh, "crop_n_layers": args.crop_n_layers, "crop_nms_thresh": args.crop_nms_thresh, "crop_overlap_ratio": args.crop_overlap_ratio, "crop_n_points_downscale_factor": args.crop_n_points_downscale_factor, "min_mask_region_area": args.min_mask_region_area, } amg_kwargs = {k: v for k, v in amg_kwargs.items() if v is not None} return amg_kwargs def main(args: argparse.Namespace) -> None: print("Loading model...") sam = sam_model_registry[args.model_type](checkpoint=args.checkpoint) _ = sam.to(device=args.device) output_mode = "coco_rle" if args.convert_to_rle else "binary_mask" amg_kwargs = get_amg_kwargs(args) generator = SamAutomaticMaskGenerator(sam, output_mode=output_mode, **amg_kwargs) if not os.path.isdir(args.input): targets = [args.input] else: targets = [ f for f in os.listdir(args.input) if not os.path.isdir(os.path.join(args.input, f)) ] targets = [os.path.join(args.input, f) for f in targets] os.makedirs(args.output, exist_ok=True) for t in targets: print(f"Processing '{t}'...") image = cv2.imread(t) if image is None: print(f"Could not load '{t}' as an image, skipping...") continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) masks = generator.generate(image) base = os.path.basename(t) base = os.path.splitext(base)[0] save_base = os.path.join(args.output, base) if output_mode == "binary_mask": os.makedirs(save_base, exist_ok=False) write_masks_to_folder(masks, save_base) else: save_file = save_base + ".json" with open(save_file, "w") as f: json.dump(masks, f) print("Done!") if __name__ == "__main__": args = parser.parse_args() main(args)