Upload HuggingFaceDemo.py

HuggingFaceDemo.py

@@ -0,0 +1,108 @@
+import os
+import sys
+
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+
+
+# Hugging Face repository
+repo_id = "ConservationDrones/DroneMegaDetector"
+
+# Load the demo image, or an image supplied on the command line.
+image_path = sys.argv[1] if len(sys.argv) > 1 else "example.png"
+image = Image.open(image_path).convert("RGB")
+
+
+# Load the processor and model from Hugging Face.
+processor = AutoImageProcessor.from_pretrained(
+    repo_id,
+)
+
+model = AutoModelForObjectDetection.from_pretrained(
+    repo_id,
+).eval()
+
+
+# Use 4 CPU threads, as in the original local demo.
+torch.set_num_threads(4)
+
+
+# Preprocess the image and run inference.
+encoded = processor(
+    images=image,
+    return_tensors="pt",
+)
+
+with torch.inference_mode():
+    outputs = model(**encoded)
+
+
+# Keep scores above 0.3 and map boxes back to the original image dimensions.
+results = processor.post_process_object_detection(
+    outputs,
+    threshold=0.3,
+    target_sizes=[(image.height, image.width)],
+)[0]
+
+
+# Print the class, confidence, and [xmin, ymin, xmax, ymax] in source pixels.
+print(
+    f"{image_path}: {len(results['scores'])} detections (score >= 0.3)"
+)
+
+for score, label, box in zip(
+    results["scores"],
+    results["labels"],
+    results["boxes"],
+):
+    print(
+        model.config.id2label[label.item()],
+        f"{score.item():.3f}",
+        [round(v, 1) for v in box.tolist()],
+    )
+
+
+# Draw the detections.
+draw = ImageDraw.Draw(image)
+font = ImageFont.load_default(size=18)
+
+for score, label, box in zip(
+    results["scores"],
+    results["labels"],
+    results["boxes"],
+):
+    bounds = box.tolist()
+
+    text = (
+        f"{model.config.id2label[label.item()]} "
+        f"{score.item():.2f}"
+    )
+
+    draw.rectangle(
+        bounds,
+        outline="red",
+        width=3,
+    )
+
+    text_x = min(
+        max(2, bounds[0]),
+        image.width - draw.textlength(text, font=font) - 2,
+    )
+
+    draw.text(
+        (text_x, max(0, bounds[1] - 20)),
+        text,
+        font=font,
+        fill="white",
+        stroke_width=2,
+        stroke_fill="black",
+    )
+
+
+# Save the annotated image.
+output_path = "example_annotated.jpg"
+image.save(output_path, quality=95)
+
+print(f"Saved {output_path}")
+