Pneumonia_obj / app.py
0llheaven's picture
Update app.py
9b8311b verified
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)