Create app.py

app.py

@@ -0,0 +1,77 @@
+import gradio as gr
+import torch
+import cv2
+import numpy as np
+from ultralytics import YOLO
+from PIL import Image
+import os
+
+# Load the YOLOv8 model ONCE (faster processing)
+MODEL_PATH = "yolov8l.pt"
+model = YOLO(MODEL_PATH)
+
+# Get sample images from Space directory
+valid_extensions = (".jpg", ".jpeg", ".png")
+preloaded_images = [img for img in os.listdir() if img.lower().endswith(valid_extensions)]
+
+# Object detection function
+def predict(image):
+    if isinstance(image, str):  # Sample image selected
+        image = Image.open(image)
+    else:  # Uploaded image
+        image = Image.fromarray(image)
+
+    results = model(image)  # YOLO detection
+
+    # Convert image for OpenCV processing
+    image_cv = np.array(image)
+    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
+    overlay = image_cv.copy()
+
+    for result in results:
+        for box in result.boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0])
+            label = model.names[int(box.cls)]
+            confidence = float(box.conf)
+
+            # Translucent blue overlay
+            cv2.rectangle(overlay, (x1, y1), (x2, y2), (255, 0, 0), -1)
+            cv2.rectangle(image_cv, (x1, y1), (x2, y2), (255, 0, 0), 2)
+            cv2.putText(image_cv, f"{label}: {confidence:.2f}", (x1, y1 - 10),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+
+    # Blend overlay with transparency
+    alpha = 0.4
+    image_masked = cv2.addWeighted(overlay, alpha, image_cv, 1 - alpha, 0)
+
+    return Image.fromarray(cv2.cvtColor(image_masked, cv2.COLOR_BGR2RGB))
+
+# Processing function with correct priority
+def process_image(image, sample_name):
+    if image is not None:  # Prioritize uploaded image
+        return predict(image)
+    elif sample_name:  # Otherwise, use the selected sample
+        return predict(sample_name)
+    return None  # No input provided
+
+# Gradio Interface
+with gr.Blocks() as interface:
+    gr.Markdown("# 🪨 Moon Rock Detection")
+    gr.Markdown("Upload a moon surface image or select a sample.")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="numpy", label="Upload Image")
+            sample_dropdown = gr.Dropdown(
+                choices=preloaded_images, label="Or Select a Sample Image", interactive=True
+            )
+
+        with gr.Column(scale=2):
+            output_image = gr.Image(type="pil", label="Detection Result")
+
+    # Automatically detect when an image is uploaded or selected
+    image_input.change(process_image, inputs=[image_input, sample_dropdown], outputs=output_image)
+    sample_dropdown.change(process_image, inputs=[image_input, sample_dropdown], outputs=output_image)
+
+if __name__ == "__main__":
+    interface.launch()