Update app.py

app.py

@@ -4,15 +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 YOLOv10n model
+# Load YOLOv8 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:
@@ -27,73 +26,240 @@ def predict_image(input_image, name, patient_id):
     elif image_np.shape[2] == 4:  # RGBA to RGB
         image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
 
-    # Perform inference with YOLOv10n model
+    # 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:
-        # Iterate through each detected object
-        for i in range(len(results[0].boxes)):
-            box = results[0].boxes[i]
-            predicted_class = int(box.cls.item())
+        for box in results[0].boxes:
+            # Get class index and confidence for each detection
+            class_index = 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
+            # 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)
 
-    return pil_image_with_boxes, "\n".join(raw_predictions)
+    # 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."
 
-    # Run prediction
     output_image, raw_result = predict_image(input_image, name, patient_id)
-    
-    return output_image, raw_result
+    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)
 
-# 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.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.Column():
-        submit_btn = gr.Button("Submit")
-        output_image = gr.Image(type="pil", label="Predicted Image")
+        with gr.Tab("Download Results"):
+            download_choice = gr.Radio(["Database (.db)", "Predicted Image (.png)"], label="Download Option")
+            download_btn = gr.Button("Download")
 
-    with gr.Row():
-        raw_result = gr.Textbox(label="Raw Result", lines=5)
+            download_output = gr.File()
 
-    submit_btn.click(fn=interface, inputs=[name, patient_id, input_image], outputs=[output_image, raw_result])
+            download_btn.click(fn=download_interface, inputs=download_choice, outputs=download_output)
 
-# Launch the Gradio app
 demo.launch()