| # Initialize the configuration | |
| cfg = get_cfg() | |
| # Load the config file from the model zoo | |
| cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")) | |
| # Set the pre-trained model weights | |
| cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml") | |
| # Set the confidence threshold for predictions | |
| cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # You can adjust this threshold | |
| # Specify the device to run on (GPU if available, else CPU) | |
| cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Create the predictor | |
| predictor = DefaultPredictor(cfg) | |
| # Path to your image | |
| image_path = "path_to_your_image.jpg" | |
| # Read the image using OpenCV | |
| image = cv2.imread(image_path) | |
| # Check if the image was loaded successfully | |
| if image is None: | |
| raise ValueError(f"Image not found at {image_path}") | |
| # Perform inference | |
| outputs = predictor(image) | |
| # Convert the image from BGR to RGB for visualization | |
| image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # Create a Visualizer instance | |
| v = Visualizer(image_rgb, metadata=model_zoo.get_cfg().MODEL.META_ARCHITECTURE, scale=1.2, instance_mode=ColorMode.IMAGE_BW) | |
| # Draw the predictions on the image | |
| out = v.draw_instance_predictions(outputs["instances"].to("cpu")) | |
| # Convert back to BGR for OpenCV compatibility (if needed) | |
| output_image = out.get_image()[:, :, ::-1] | |
| # Display the image using OpenCV | |
| cv2.imshow("Object Detection", output_image) | |
| cv2.waitKey(0) # Press any key to close the window | |
| cv2.destroyAllWindows() | |
| # Alternatively, display using matplotlib | |
| plt.figure(figsize=(12, 8)) | |
| plt.imshow(out.get_image()) | |
| plt.axis('off') | |
| plt.title("Object Detection Results") | |
| plt.show() | |