import gradio as gr from ultralytics import YOLO import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont import base64 from io import BytesIO import zipfile import os from pathlib import Path # Load YOLOv8 model model = YOLO("best.pt") # Define paths for uploaded and predicted images uploaded_folder = Path('Uploaded_Picture') predicted_folder = Path('Predicted_Picture') uploaded_folder.mkdir(parents=True, exist_ok=True) predicted_folder.mkdir(parents=True, exist_ok=True) # Path for HTML database file html_db_file = Path('patient_predictions.html') # Initialize HTML file if not present if not html_db_file.exists(): with open(html_db_file, 'w') as f: f.write(""" Patient Prediction Database

Patient Prediction Database

""") def predict_image(input_image, name, age, medical_record, sex): # Ensure input image is provided if input_image is None: return None, "Please upload an image for prediction." # Convert PIL image to NumPy array image_np = np.array(input_image) # Perform YOLO prediction results = model(image_np) image_with_boxes = image_np.copy() label = "Unknown" if results[0].boxes: # Take the result with the highest confidence best_result = max(results[0].boxes, key=lambda x: x.conf.item()) class_index = best_result.cls.item() # Determine class label if class_index == 0: label = "Immature" color = (0, 255, 255) elif class_index == 1: label = "Mature" color = (255, 0, 0) else: label = "Normal" color = (0, 255, 0) confidence = best_result.conf.item() xmin, ymin, xmax, ymax = map(int, best_result.xyxy[0]) # Draw bounding box and label on image cv2.rectangle(image_with_boxes, (xmin, ymin), (xmax, ymax), color, 2) font_scale, thickness = 1.0, 2 cv2.putText(image_with_boxes, f'{label} {confidence:.2f}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, thickness) # Convert the annotated image back to PIL pil_image_with_boxes = Image.fromarray(image_with_boxes) # Save images to folders image_name = f"{name}_{age}_{medical_record}_{sex}.png" input_image.save(uploaded_folder / image_name) pil_image_with_boxes.save(predicted_folder / image_name) # Convert predicted image to base64 for embedding in HTML buffered = BytesIO() pil_image_with_boxes.save(buffered, format="PNG") predicted_image_base64 = base64.b64encode(buffered.getvalue()).decode() # Append patient information to HTML append_patient_info_to_html(name, age, medical_record, sex, label, predicted_image_base64) raw_prediction = f"Name: {name}, Age: {age}, Medical Record: {medical_record}, Sex: {sex}, Result: {label}" return pil_image_with_boxes, raw_prediction def append_patient_info_to_html(name, age, medical_record, sex, result, predicted_image_base64): # Append a new patient entry to the HTML file html_entry = f""" """ with open(html_db_file, 'a') as f: f.write(html_entry) # Close the HTML file after writing (for proper structure) with open(html_db_file, 'a') as f: f.write("
Name Age Medical Record Sex Result Predicted Image
{name} {age} {medical_record} {sex} {result} Predicted Image
") return str(html_db_file) def download_uploaded_folder(): # Create a zip file of the uploaded folder zip_path = 'uploaded_images.zip' with zipfile.ZipFile(zip_path, 'w') as zf: for file in uploaded_folder.iterdir(): zf.write(file, arcname=file.name) return zip_path def download_predicted_folder(): # Create a zip file of the predicted folder zip_path = 'predicted_images.zip' with zipfile.ZipFile(zip_path, 'w') as zf: for file in predicted_folder.iterdir(): zf.write(file, arcname=file.name) return zip_path # Launch Gradio Interface 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.") gr.Markdown("This application uses YOLOv8 with mAP=0.981") with gr.Column(): name = gr.Textbox(label="Name") age = gr.Number(label="Age") medical_record = gr.Number(label="Medical Record") sex = gr.Radio(["Male", "Female"], label="Sex") 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="Prediction Result") with gr.Row(): download_html_btn = gr.Button("Download Patient Information (HTML)") download_uploaded_btn = gr.Button("Download Uploaded Images") download_predicted_btn = gr.Button("Download Predicted Images") # Add file download output components for the uploaded and predicted images patient_info_file = gr.File(label="Patient Information HTML File") uploaded_folder_file = gr.File(label="Uploaded Images Zip File") predicted_folder_file = gr.File(label="Predicted Images Zip File") # Connect functions with components submit_btn.click(fn=predict_image, inputs=[name, age, medical_record, sex, input_image], outputs=[output_image, raw_result]) download_html_btn.click(fn=append_patient_info_to_html, inputs=[name, age, medical_record, sex, raw_result], outputs=patient_info_file) download_uploaded_btn.click(fn=download_uploaded_folder, outputs=uploaded_folder_file) download_predicted_btn.click(fn=download_predicted_folder, outputs=predicted_folder_file) # Launch Gradio app demo.launch()