Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import cv2 | |
| from PIL import Image | |
| import numpy as np | |
| from ultralytics import YOLO | |
| model = YOLO('best_V4.pt') | |
| def predict(image): | |
| results = model(image, conf=0.8) | |
| detected = False | |
| LABEL_MAP = { | |
| 0: "Other", | |
| 1: "Pneumonia" | |
| } | |
| labels_found = [] | |
| for result in results: | |
| boxes = result.boxes.xyxy.cpu().numpy() | |
| confidences = result.boxes.conf.cpu().numpy() | |
| class_ids = result.boxes.cls.cpu().numpy() | |
| for box, confidence, class_id in zip(boxes, confidences, class_ids): | |
| x1, y1, x2, y2 = map(int, box[:4]) | |
| label = LABEL_MAP.get(int(class_id), "Unknown") | |
| cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| label_text = f"{label} {confidence:.2f}" | |
| labels_found.append(label_text) | |
| cv2.putText(image, label_text, (x1, y1 - 10), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) | |
| detected = True | |
| if not detected: | |
| return None, "No detected" | |
| pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
| message = "\n".join(labels_found) | |
| return pil_image, message | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="numpy"), | |
| outputs=[ | |
| gr.Image(type="pil", label="Detection Result"), | |
| gr.Textbox(label="Message") | |
| ], | |
| allow_flagging="never", | |
| api_name=False # <- THIS DISABLES OPENAPI GENERATION AND AVOIDS THE ERROR | |
| ) | |
| demo.launch(share=True, debug=True) |