File size: 1,587 Bytes
472d78b |
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 |
import torch
import detectron2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2
import json
import argparse
def main():
parser = argparse.ArgumentParser(description="Run inference with Detectron2 model")
parser.add_argument("--image", required=True, help="Path to input image")
parser.add_argument("--output", default="output.jpg", help="Path to output image")
args = parser.parse_args()
# Load config
cfg = get_cfg()
with open("config.json", "r") as f:
cfg_dict = json.load(f)
cfg.merge_from_dict(cfg_dict)
# Update inference parameters
cfg.MODEL.WEIGHTS = "model.pth"
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Create predictor
predictor = DefaultPredictor(cfg)
# Load image
image = cv2.imread(args.image)
# Run prediction
outputs = predictor(image)
# Load metadata
with open("metadata.json", "r") as f:
metadata_dict = json.load(f)
# Setup metadata
metadata = MetadataCatalog.get("inference")
metadata.thing_classes = metadata_dict["thing_classes"]
# Visualize
v = Visualizer(image[:, :, ::-1], metadata=metadata, scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# Save output
cv2.imwrite(args.output, out.get_image()[:, :, ::-1])
print(f"Saved output to {args.output}")
if __name__ == "__main__":
main()
|