import gradio as gr from ultralytics import YOLO import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont import sqlite3 import base64 from io import BytesIO import tempfile import pandas as pd # Load YOLOv8 model model = YOLO("best.pt") # 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 prediction results = model(image_np) # Draw bounding boxes on the image image_with_boxes = image_np.copy() raw_predictions = [] label = "Unknown" # Default label if no detection if results[0].boxes: for box in results[0].boxes: # Get class index and confidence for each detection class_index = box.cls.item() confidence = box.conf.item() # Determine the label based on the class index if class_index == 0: label = "Mature" color = (255, 0, 0) # Red for Mature elif class_index == 1: label = "Immature" color = (0, 255, 255) # Yellow for Immature else: label = "Normal" color = (0, 255, 0) # Green for Normal xmin, ymin, xmax, ymax = map(int, box.xyxy[0]) # Draw the bounding box cv2.rectangle(image_with_boxes, (xmin, ymin), (xmax, ymax), color, 2) # Enlarge font scale and thickness font_scale = 1.0 thickness = 2 # Calculate label background size (text_width, text_height), baseline = cv2.getTextSize(f'{label} {confidence:.2f}', cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness) cv2.rectangle(image_with_boxes, (xmin, ymin - text_height - baseline), (xmin + text_width, ymin), (0, 0, 0), cv2.FILLED) # Put the label text with black background cv2.putText(image_with_boxes, f'{label} {confidence:.2f}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness) raw_predictions.append(f"Label: {label}, Confidence: {confidence:.2f}, Box: [{xmin}, {ymin}, {xmax}, {ymax}]") raw_predictions_str = "\n".join(raw_predictions) # Convert to PIL image for further processing pil_image_with_boxes = Image.fromarray(image_with_boxes) # Add text and watermark pil_image_with_boxes = add_text_and_watermark(pil_image_with_boxes, name, patient_id, label) return pil_image_with_boxes, raw_predictions_str # Function to add watermark def add_watermark(image): try: logo = Image.open('image-logo.png').convert("RGBA") image = image.convert("RGBA") # Resize logo basewidth = 100 wpercent = (basewidth / float(logo.size[0])) hsize = int((float(wpercent) * logo.size[1])) logo = logo.resize((basewidth, hsize), Image.LANCZOS) # Position logo position = (image.width - logo.width - 10, image.height - logo.height - 10) # Composite image transparent = Image.new('RGBA', (image.width, image.height), (0, 0, 0, 0)) transparent.paste(image, (0, 0)) transparent.paste(logo, position, mask=logo) return transparent.convert("RGB") except Exception as e: print(f"Error adding watermark: {e}") return image # Function to add text and watermark def add_text_and_watermark(image, name, patient_id, label): draw = ImageDraw.Draw(image) # Load a larger font (adjust the size as needed) font_size = 48 # Example font size try: font = ImageFont.truetype("font.ttf", size=font_size) except IOError: font = ImageFont.load_default() print("Error: cannot open resource, using default font.") text = f"Name: {name}, ID: {patient_id}, Result: {label}" # Calculate text bounding box text_bbox = draw.textbbox((0, 0), text, font=font) text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1] text_x = 20 text_y = 40 padding = 10 # Draw a filled rectangle for the background draw.rectangle( [text_x - padding, text_y - padding, text_x + text_width + padding, text_y + text_height + padding], fill="black" ) # Draw text on top of the rectangle draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=font) # Add watermark to the image image_with_watermark = add_watermark(image) return image_with_watermark # Function to initialize the database def init_db(): conn = sqlite3.connect('results.db') c = conn.cursor() c.execute('''CREATE TABLE IF NOT EXISTS results (id INTEGER PRIMARY KEY, name TEXT, patient_id TEXT, input_image BLOB, predicted_image BLOB, result TEXT)''') conn.commit() conn.close() # Function to submit result to the database def submit_result(name, patient_id, input_image, predicted_image, result): conn = sqlite3.connect('results.db') c = conn.cursor() input_image_np = np.array(input_image) _, input_buffer = cv2.imencode('.png', cv2.cvtColor(input_image_np, cv2.COLOR_RGB2BGR)) input_image_bytes = input_buffer.tobytes() predicted_image_np = np.array(predicted_image) predicted_image_rgb = cv2.cvtColor(predicted_image_np, cv2.COLOR_RGB2BGR) # Ensure correct color conversion _, predicted_buffer = cv2.imencode('.png', predicted_image_rgb) predicted_image_bytes = predicted_buffer.tobytes() c.execute("INSERT INTO results (name, patient_id, input_image, predicted_image, result) VALUES (?, ?, ?, ?, ?)", (name, patient_id, input_image_bytes, predicted_image_bytes, result)) conn.commit() conn.close() return "Result submitted to database." # Function to load and view database def view_database(): conn = sqlite3.connect('results.db') c = conn.cursor() c.execute("SELECT * FROM results") rows = c.fetchall() conn.close() # Convert to pandas DataFrame df = pd.DataFrame(rows, columns=["ID", "Name", "Patient ID", "Input Image", "Predicted Image", "Result"]) return df # Function to download database or image def download_file(choice): conn = sqlite3.connect('results.db') c = conn.cursor() if choice == "Database (.db)": conn.close() return 'results.db' else: c.execute("SELECT predicted_image FROM results ORDER BY id DESC LIMIT 1") row = c.fetchone() conn.close() if row: image_bytes = row[0] with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file: temp_file.write(image_bytes) temp_file.flush() # Ensure all data is written before closing return temp_file.name else: conn.close() raise FileNotFoundError("No images found in the database.") # Initialize the database init_db() # Gradio Interface def interface(name, patient_id, input_image): if input_image is None: return "Please upload an image." output_image, raw_result = predict_image(input_image, name, patient_id) submit_status = submit_result(name, patient_id, input_image, output_image, raw_result) return output_image, raw_result, submit_status # View Database Function def view_db_interface(): df = view_database() return df # Download Function def download_interface(choice): try: file_path = download_file(choice) with open(file_path, "rb") as file: return file.read(), f"{choice}" except Exception as e: return str(e) with gr.Blocks() as demo: with gr.Tabs(): with gr.Tab("Image Analyzer and Screener"): gr.Markdown("## Cataract Detection System") with gr.Row(): with gr.Column(): input_image = gr.Image(label="Upload Image") name = gr.Textbox(label="Patient Name") patient_id = gr.Textbox(label="Patient ID") submit_btn = gr.Button("Submit") with gr.Column(): output_image = gr.Image(label="Predicted Image") raw_result = gr.Textbox(label="Raw Result") submit_status = gr.Textbox(label="Submission Status") submit_btn.click(fn=interface, inputs=[name, patient_id, input_image], outputs=[output_image, raw_result, submit_status]) with gr.Tab("Database Viewer"): view_db_btn = gr.Button("View Database") database_display = gr.Dataframe() view_db_btn.click(fn=view_db_interface, outputs=database_display) with gr.Tab("Download Results"): download_choice = gr.Radio(["Database (.db)", "Predicted Image (.png)"], label="Download Option") download_btn = gr.Button("Download") download_output = gr.File() download_btn.click(fn=download_interface, inputs=download_choice, outputs=download_output) demo.launch()