File size: 2,721 Bytes
c2efdf5 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
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()
|