File size: 1,708 Bytes
7a0e0e7 8f1e0df fcddb1d 7a0e0e7 |
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 |
# 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()
|