Create app.py

app.py

@@ -0,0 +1,53 @@
+import gradio as gr
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+import torch
+from PIL import Image, ImageDraw
+
+# Load the model and processor
+processor = AutoImageProcessor.from_pretrained("0llheaven/CON-DETR-V5")
+model = AutoModelForObjectDetection.from_pretrained("0llheaven/CON-DETR-V5")
+
+def detect_objects(image, score_threshold):
+    # Convert image to RGB if it's grayscale
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    
+    # Prepare input for the model
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # Filter predictions based on the user-defined score threshold
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)
+
+    labels_output = [] 
+    
+    # Draw bounding boxes around detected objects
+    draw = ImageDraw.Draw(image)
+    for result in results:
+        scores = result["scores"]
+        labels = result["labels"]
+        boxes = result["boxes"]
+
+        for score, label, box in zip(scores, labels, boxes):
+            if score >= score_threshold:  # Only draw if score is above threshold
+                box = [round(i, 2) for i in box.tolist()]
+                label_name = "Pneumonia" if label.item() == 0 else "Other"
+                draw.rectangle(box, outline="red", width=3)
+                draw.text((box[0], box[1]), f"{label_name}: {round(score.item(), 3)}", fill="red")
+                labels_output.append(f"{label_name}: {round(score.item(), 3)}")
+    
+    return image, "\n".join(labels_output)
+
+# Create the Gradio interface
+interface = gr.Interface(
+    fn=detect_objects, 
+    inputs=[gr.Image(type="pil"), gr.Slider(0, 1, value=0.5, label="Score Threshold")],  # Add slider for score threshold
+    # outputs=gr.Image(type="pil"),  # Corrected output type
+    outputs=[gr.Image(type="pil"), gr.Textbox(label="Detected Objects")],
+    title="Object Detection with Transformers",
+    description="Upload an image to detect objects using a fine-tuned Conditional-DETR model."
+)
+
+# Launch the interface
+interface.launch()