Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| from PIL import Image, ImageOps, ImageEnhance | |
| import numpy as np | |
| import tempfile | |
| # Load your model | |
| model = YOLO("model/best.pt") | |
| def preprocess(image): | |
| """Safe preprocessing for PIL or numpy input.""" | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| image = ImageOps.exif_transpose(image).convert("RGB") | |
| # Optional resize for performance | |
| w, h = image.size | |
| max_dim = max(w, h) | |
| if max_dim > 1024: | |
| scale = 1024 / max_dim | |
| image = image.resize((int(w * scale), int(h * scale)), Image.LANCZOS) | |
| # Light contrast enhancement | |
| image = ImageEnhance.Contrast(image).enhance(1.05) | |
| return image | |
| def detect(image, conf=0.5, iou=0.5): | |
| """Run YOLO detection on a single model Space.""" | |
| image = preprocess(image) | |
| results = model.predict(image, conf=conf, iou=iou) | |
| boxes = results[0].boxes | |
| # Convert YOLO output to numpy RGB | |
| output = results[0].plot()[:, :, ::-1] # BGR → RGB | |
| if len(boxes) > 0: | |
| diagnosis = "⚠️ Bleeding detected." | |
| else: | |
| diagnosis = "🟢 No bleeding detected." | |
| return [output, diagnosis] | |
| # Gradio Interface | |
| interface = gr.Interface( | |
| fn=detect, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload Image"), | |
| gr.Slider(0, 1, value=0.5, step=0.05, label="Confidence Threshold"), | |
| gr.Slider(0, 1, value=0.5, step=0.05, label="NMS IoU Threshold"), | |
| ], | |
| outputs=[ | |
| gr.Image(label="Bleeding Detection Result"), | |
| gr.Textbox(label="Diagnosis") | |
| ], | |
| title="Bleeding Detection" | |
| ) | |
| interface.launch() | |