Create owl.py

owl.py

@@ -0,0 +1,84 @@
+import requests
+from PIL import Image, ImageDraw, ImageFont
+import torch
+import os
+
+from transformers import Owlv2Processor, Owlv2ForObjectDetection
+
+# Load the model and processor
+processor = Owlv2Processor.from_pretrained("google/owlv2-large-patch14-ensemble")
+model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-large-patch14-ensemble")
+
+# Option 1: Load image from local file
+image_path = "image.jpg"  # Replace with your image path
+image = Image.open(image_path)
+
+
+# Define what you want to detect
+text_labels = [["a person with a hat"]]
+
+# Process the image and text
+inputs = processor(text=text_labels, images=image, return_tensors="pt")
+outputs = model(**inputs)
+
+# Target image sizes (height, width) to rescale box predictions [batch_size, 2]
+target_sizes = torch.tensor([(image.height, image.width)])
+
+# Convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
+results = processor.post_process_grounded_object_detection(
+    outputs=outputs, target_sizes=target_sizes, threshold=0.1, text_labels=text_labels
+)
+
+# Retrieve predictions for the first image
+result = results[0]
+boxes, scores, text_labels_detected = result["boxes"], result["scores"], result["text_labels"]
+
+# Create a copy of the original image for drawing
+output_image = image.copy()
+draw = ImageDraw.Draw(output_image)
+
+# Try to use a default font, fallback to default if not available
+try:
+    font = ImageFont.truetype("arial.ttf", 16)
+except OSError:
+    font = ImageFont.load_default()
+
+# Colors for different detections
+colors = ["red", "blue", "green", "orange", "purple", "yellow", "pink", "cyan"]
+
+print("Detection Results:")
+print("-" * 50)
+
+# Draw bounding boxes and labels
+for i, (box, score, text_label) in enumerate(zip(boxes, scores, text_labels_detected)):
+    box = [round(i, 2) for i in box.tolist()]
+    confidence = round(score.item(), 3)
+    
+    print(f"Detected {text_label} with confidence {confidence} at location {box}")
+    
+    # Get coordinates
+    xmin, ymin, xmax, ymax = box
+    
+    # Choose color
+    color = colors[i % len(colors)]
+    
+    # Draw bounding box
+    draw.rectangle([xmin, ymin, xmax, ymax], outline=color, width=3)
+    
+    # Draw label with confidence
+    label_text = f"{text_label}: {confidence}"
+    
+    # Get text bounding box for background
+    bbox = draw.textbbox((xmin, ymin - 25), label_text, font=font)
+    
+    # Draw background rectangle for text
+    draw.rectangle([bbox[0]-2, bbox[1]-2, bbox[2]+2, bbox[3]+2], fill=color)
+    
+    # Draw text
+    draw.text((xmin, ymin - 25), label_text, fill="white", font=font)
+
+# Save the output image
+output_path = "output_img.jpg"
+output_image.save(output_path)
+print(f"\nOutput image saved as: {output_path}")
+