Update app.py

app.py

@@ -3,12 +3,10 @@ from ultralytics import YOLO
 import cv2
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
-import base64
-from io import BytesIO
-import tempfile
 import os
 from pathlib import Path
 import shutil
+import tempfile
 
 # Load YOLOv8 model
 model = YOLO("best.pt")
@@ -19,10 +17,9 @@ predicted_folder = Path('Predicted_Picture')
 uploaded_folder.mkdir(parents=True, exist_ok=True)
 predicted_folder.mkdir(parents=True, exist_ok=True)
 
-# Path to store accumulated HTML data
-html_file_path = Path(tempfile.gettempdir()) / 'patient_data.html'
+# Global patient data list to accumulate HTML data
+patient_data = []
 
-# Function to predict image and add bounding box, text, circle, and watermark
 def predict_image(input_image, name, age, medical_record, sex):
     if input_image is None:
         return None, "Please Input The Image"
@@ -39,7 +36,7 @@ def predict_image(input_image, name, age, medical_record, sex):
     # Perform prediction
     results = model(image_np)
 
-    # Draw bounding boxes on the image
+    # Draw bounding boxes and white circle on the image
     image_with_boxes = image_np.copy()
     raw_predictions = []
 
@@ -65,17 +62,13 @@ def predict_image(input_image, name, age, medical_record, sex):
         # Draw the bounding box
         cv2.rectangle(image_with_boxes, (xmin, ymin), (xmax, ymax), color, 2)
 
-        # Calculate the center of the bounding box
-        center_x = (xmin + xmax) // 2
-        center_y = (ymin + ymax) // 2
-        
-        # Calculate the radius (1/12 of the average of the width and height of the bounding box)
+        # Draw the white circle in the center of the bounding box
         box_width = xmax - xmin
         box_height = ymax - ymin
-        radius = int((box_width + box_height) / 24)  # Average of width and height divided by 12
-
-        # Draw a white circle at the center of the bounding box
-        cv2.circle(image_with_boxes, (center_x, center_y), radius, (255, 255, 255), thickness=2)
+        center_x = xmin + box_width // 2
+        center_y = ymin + box_height // 2
+        radius = int((box_width + box_height) / 2 / 12)
+        cv2.circle(image_with_boxes, (center_x, center_y), radius, (255, 255, 255), 2)
         
         # Enlarge font scale and thickness
         font_scale = 1.0
@@ -105,31 +98,6 @@ def predict_image(input_image, name, age, medical_record, sex):
     
     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, age, medical_record, sex, label):
     draw = ImageDraw.Draw(image)
@@ -160,36 +128,29 @@ def add_text_and_watermark(image, name, age, medical_record, sex, label):
     # 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
+    return image
 
 # Function to save patient info in HTML and accumulate data
 def save_patient_info_to_html(name, age, medical_record, sex, result):
+    global patient_data
+    new_data = f"<p><strong>Name:</strong> {name}, <strong>Age:</strong> {age}, <strong>Medical Record:</strong> {medical_record}, <strong>Sex:</strong> {sex}, <strong>Result:</strong> {result}</p>"
+    patient_data.append(new_data)
+
     html_content = f"""
     <html>
         <body>
             <h1>Patient Information</h1>
-            <p><strong>Name:</strong> {name}</p>
-            <p><strong>Age:</strong> {age}</p>
-            <p><strong>Medical Record:</strong> {medical_record}</p>
-            <p><strong>Sex:</strong> {sex}</p>
-            <p><strong>Result:</strong> {result}</p>
-            <hr>
+            {''.join(patient_data)}
         </body>
     </html>
     """
     
-    # Check if the HTML file already exists
-    if html_file_path.exists():
-        with open(html_file_path, 'a') as f:
-            f.write(html_content)
-    else:
-        with open(html_file_path, 'w') as f:
-            f.write(html_content)
+    # Save HTML content to file
+    html_file_path = os.path.join(tempfile.gettempdir(), 'patient_info.html')
+    with open(html_file_path, 'w') as f:
+        f.write(html_content)
     
-    return str(html_file_path)
+    return html_file_path
 
 # Function to download the folders
 def download_folder(folder):
@@ -201,28 +162,6 @@ def download_folder(folder):
     return zip_path
 
 # Gradio Interface
-def interface(name, age, medical_record, sex, input_image):
-    if input_image is None:
-        return None, "Please upload an image.", None
-
-    output_image, raw_result = predict_image(input_image, name, age, medical_record, sex)
-    
-    if output_image is None:
-        return None, raw_result, None
-    
-    # Save patient info to HTML
-    html_file_path = save_patient_info_to_html(name, age, medical_record, sex, raw_result)
-    
-    # Encode the image to display in Gradio
-    buffered = BytesIO()
-    output_image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    
-    # Provide the zip file path for download
-    zip_file = download_folder(predicted_folder)
-    
-    return f'<img src="data:image/png;base64,{img_str}" alt="Processed Image"/>', raw_result, zip_file
-
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("# Cataract Detection System")
@@ -253,10 +192,10 @@ with gr.Blocks() as demo:
     uploaded_folder_file = gr.File(label="Uploaded Images Zip File")
     predicted_folder_file = gr.File(label="Predicted Images Zip File")
 
-    submit_btn.click(fn=interface, inputs=[name, age, medical_record, sex, input_image], outputs=[output_image, raw_result])
+    submit_btn.click(fn=predict_image, inputs=[name, age, medical_record, sex, input_image], outputs=[output_image, raw_result])
     download_html_btn.click(fn=save_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)
+    download_uploaded_btn.click(fn=download_folder, inputs=[uploaded_folder], outputs=uploaded_folder_file)
+    download_predicted_btn.click(fn=download_folder, inputs=[predicted_folder], outputs=predicted_folder_file)
 
 # Launch Gradio app
 demo.launch()