Update app.py

app.py

@@ -4,13 +4,14 @@ 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"
@@ -30,51 +31,50 @@ def predict_image(input_image, name, patient_id):
     # 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 = "Immature"
-                color = (0, 255, 255)  # Yellow for Immature
-            elif class_index == 1:
-                label = "Mature"
-                color = (255, 0, 0)  # Red for Mature
-            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}]")
+        # Sort the results by confidence and take the highest confidence one
+        highest_confidence_result = max(results[0].boxes, key=lambda x: x.conf.item())
+
+        # Determine the label based on the class index
+        class_index = highest_confidence_result.cls.item()
+        if class_index == 0:
+            label = "Immature"
+            color = (0, 255, 255)  # Yellow for Immature
+        elif class_index == 1:
+            label = "Mature"
+            color = (255, 0, 0)  # Red for Mature
+        else:
+            label = "Normal"
+            color = (0, 255, 0)  # Green for Normal
+
+        confidence = highest_confidence_result.conf.item()
+        xmin, ymin, xmax, ymax = map(int, highest_confidence_result.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
@@ -82,21 +82,21 @@ 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}")
@@ -105,7 +105,7 @@ def add_watermark(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:
@@ -113,9 +113,9 @@ def add_text_and_watermark(image, name, patient_id, label):
     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]
@@ -128,13 +128,13 @@ def add_text_and_watermark(image, name, patient_id, label):
         [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
@@ -150,16 +150,16 @@ def init_db():
 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()
@@ -170,24 +170,28 @@ def submit_result(name, patient_id, input_image, predicted_image, result):
 def view_database():
     conn = sqlite3.connect('results.db')
     c = conn.cursor()
-    c.execute("SELECT * FROM results")
+    c.execute("SELECT name, patient_id 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"])
-
+    
+    # Convert to pandas DataFrame for better display in Gradio
+    df = pd.DataFrame(rows, columns=["Name", "Patient ID"])
+    
     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()
+        # Provide the path to the database file
         return 'results.db'
+    elif choice == "Database (.html)":
+        df = view_database()
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.html') as temp_file:
+            df.to_html(temp_file.name)
+            return temp_file.name
     else:
+        conn = sqlite3.connect('results.db')
+        c = conn.cursor()
         c.execute("SELECT predicted_image FROM results ORDER BY id DESC LIMIT 1")
         row = c.fetchone()
         conn.close()
@@ -198,7 +202,6 @@ def download_file(choice):
                 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
@@ -211,52 +214,37 @@ def interface(name, patient_id, input_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(), file_path
-    except Exception as e:
-        return f"Error: {str(e)}", None
-
-# Build Gradio Interface
-app = gr.Blocks()
-
-with app:
-    gr.Markdown("# Eye Condition Detection System")
-
-    with gr.Row():
-        with gr.Column():
-            name = gr.Textbox(label="Name")
-            patient_id = gr.Textbox(label="Patient ID")
-            input_image = gr.Image(label="Input Image", tool="editor", type="pil")
-
-        with gr.Column():
-            output_image = gr.Image(label="Predicted Image")
-            raw_result = gr.Textbox(label="Raw Predictions", lines=5)
-            submit_status = gr.Textbox(label="Submit Status")
-
-        predict_button = gr.Button("Predict")
-
-        predict_button.click(fn=interface, inputs=[name, patient_id, input_image], outputs=[output_image, raw_result, submit_status])
-
-    with gr.Row():
-        with gr.Column():
-            view_button = gr.Button("View Database")
-            download_choice = gr.Dropdown(label="Download Option", choices=["Database (.db)", "Predicted Image (.png)"])
-            download_button = gr.Button("Download")
-
-            view_button.click(fn=view_db_interface, inputs=[], outputs=[gr.Dataframe()])
-            download_button.click(fn=download_interface, inputs=[download_choice], outputs=[gr.File(), gr.Textbox()])
-
-# Launch the Gradio app
-app.launch()
+    return output_image, submit_status
+
+inputs = [
+    gr.Textbox(label="Name"),
+    gr.Textbox(label="Patient ID"),
+    gr.Image(type="pil", label="Input Image")
+]
+
+outputs = [
+    gr.Image(label="Output Image"),
+    gr.Textbox(label="Status")
+]
+
+# File download interface
+download_inputs = gr.Radio(["Database (.db)", "Database (.html)", "Image (.png)"], label="Download Type")
+download_output = gr.File(label="Download File")
+
+app = gr.Interface(
+    fn=interface,
+    inputs=inputs,
+    outputs=outputs,
+    title="AI Cataract Detector",
+    description="Upload an image, enter the patient's name and ID, and receive a prediction."
+)
+
+download_app = gr.Interface(
+    fn=download_file,
+    inputs=download_inputs,
+    outputs=download_output,
+    title="Download Results"
+)
+
+# Combine both interfaces in one layout
+gr.TabbedInterface([app, download_app], ["Prediction", "Download"]).launch()