Spaces:
Sleeping
Sleeping
| 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() |