# 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()