File size: 2,785 Bytes
98405c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import argparse
import logging
import torch
import os
import sys
# Ensure repository root is available on sys.path when executed as a script.
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
sys.path.insert(0, PROJECT_ROOT)
from demo.inference_utils import (
build_inference_inputs,
get_metadata,
load_model,
prepare_image_inputs,
resolve_category_ids,
setup_cfg,
)
from utils.visualizer import SegmentationResultVisualizer
def parse_args():
parser = argparse.ArgumentParser(description="OpenWorldSAM2 Instance Segmentation Inference")
parser.add_argument("--config-file", required=True, help="Path to the config file")
parser.add_argument("--image", required=True, help="Path to the input image")
parser.add_argument(
"--prompts",
required=True,
nargs="+",
help="List of textual prompts describing the desired instance categories",
)
parser.add_argument("--weights", default=None, help="Path to model weights")
parser.add_argument(
"--device",
default="cuda" if torch.cuda.is_available() else "cpu",
help="Computation device",
)
parser.add_argument("--output", default="outputs/instance_result.png", help="Path to save the visualization")
parser.add_argument("--opts", default=None, nargs=argparse.REMAINDER, help="Additional config options")
return parser.parse_args()
def main():
args = parse_args()
logging.basicConfig(level=logging.INFO)
cfg = setup_cfg(args.config_file, weights=args.weights, device=args.device, opts=args.opts)
cfg.MODEL.OpenWorldSAM2.TEST.INSTANCE_ON = True
cfg.MODEL.OpenWorldSAM2.TEST.SEMANTIC_ON = False
cfg.MODEL.OpenWorldSAM2.TEST.PANOPTIC_ON = False
cfg.MODEL.OpenWorldSAM2.TEST.REFER_ON = False
# adjusting post-processing thresholds for instance segmentation
cfg.MODEL.OpenWorldSAM2.TEST.NMS_THRESHOLD = 0.2
cfg.MODEL.OpenWorldSAM2.TEST.IOU_THRESHOLD = 0.9
metadata = get_metadata(cfg)
prompts = [p.strip() for p in args.prompts]
category_ids = resolve_category_ids(prompts, metadata)
model = load_model(cfg)
image_bgr, sam_tensor, beit_tensor, height, width = prepare_image_inputs(args.image, cfg.INPUT.FORMAT)
inputs = build_inference_inputs(sam_tensor, beit_tensor, height, width, prompts, category_ids)
with torch.no_grad():
outputs = model(inputs)[0]
instances = outputs.get("instances")
visualizer = SegmentationResultVisualizer(metadata=metadata, input_format=cfg.INPUT.FORMAT)
visualizer.save_instance_result(image_bgr, instances, args.output)
logging.info("Saved instance segmentation result to %s", args.output)
if __name__ == "__main__":
main()
|