Spaces:
Sleeping
Sleeping
File size: 3,481 Bytes
4cd3714 |
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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import sqlite3
import pandas as pd
# Load YOLOv10n model
model = YOLO("best.pt")
# Define label mappings
label_mapping = {0: 'immature', 1: 'mature', 2: 'normal'}
inverse_label_mapping = {'immature': 0, 'mature': 1, 'normal': 2}
# Function to perform prediction
def predict_image(input_image, name, patient_id):
if input_image is None:
return None, "Please Input The Image"
# Convert Gradio input image (PIL Image) to numpy array
image_np = np.array(input_image)
# Ensure the image is in the correct format
if len(image_np.shape) == 2: # grayscale to RGB
image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
elif image_np.shape[2] == 4: # RGBA to RGB
image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
# Perform inference with YOLOv10n model
results = model(image_np)
# Draw bounding boxes on the image
image_with_boxes = image_np.copy()
raw_predictions = []
if results[0].boxes:
# Iterate through each detected object
for i in range(len(results[0].boxes)):
box = results[0].boxes[i]
predicted_class = int(box.cls.item())
confidence = box.conf.item()
# Apply confidence threshold
if confidence >= 0.5:
# Map the predicted class to the label
label = label_mapping[predicted_class]
# Get the bounding box coordinates
xmin, ymin, xmax, ymax = map(int, box.xyxy[0])
# Assign color for the label
color = (0, 255, 0) if label == 'normal' else (0, 255, 255) if label == 'immature' else (255, 0, 0)
# Draw the bounding box
cv2.rectangle(image_with_boxes, (xmin, ymin), (xmax, ymax), color, 2)
# Draw the label with confidence
cv2.putText(image_with_boxes, f'{label} {confidence:.2f}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
raw_predictions.append(f"Label: {label}, Confidence: {confidence:.2f}, Box: [{xmin}, {ymin}, {xmax}, {ymax}]")
# Convert to PIL image for Gradio output
pil_image_with_boxes = Image.fromarray(image_with_boxes)
return pil_image_with_boxes, "\n".join(raw_predictions)
# Gradio Interface
def interface(name, patient_id, input_image):
if input_image is None:
return "Please upload an image."
# Run prediction
output_image, raw_result = predict_image(input_image, name, patient_id)
return output_image, raw_result
# Gradio Blocks
with gr.Blocks() as demo:
with gr.Column():
gr.Markdown("# Cataract Detection System")
gr.Markdown("Upload an image to detect cataract and add patient details.")
with gr.Column():
name = gr.Textbox(label="Name")
patient_id = gr.Textbox(label="Patient ID")
input_image = gr.Image(type="pil", label="Upload an Image", image_mode="RGB")
with gr.Column():
submit_btn = gr.Button("Submit")
output_image = gr.Image(type="pil", label="Predicted Image")
with gr.Row():
raw_result = gr.Textbox(label="Raw Result", lines=5)
submit_btn.click(fn=interface, inputs=[name, patient_id, input_image], outputs=[output_image, raw_result])
# Launch the Gradio app
demo.launch() |