Spaces:

WellnessNexus
/

ThesisCataractLatest

Sleeping

App Files Files Community

ariankhalfani commited on Sep 14

Commit

22a68d2

•

1 Parent(s): 822a6f7

Update app2.py

Browse files

Files changed (1) hide show

app2.py +113 -48

app2.py CHANGED Viewed

@@ -3,10 +3,12 @@ from ultralytics import YOLO
 import cv2
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
 import os
 from pathlib import Path
 import shutil
-import tempfile
 # Load YOLOv8 model
 model = YOLO("best.pt")
@@ -17,8 +19,13 @@ predicted_folder = Path('Predicted_Picture')
 uploaded_folder.mkdir(parents=True, exist_ok=True)
 predicted_folder.mkdir(parents=True, exist_ok=True)
-# Global patient data list to accumulate HTML data
-patient_data = []
 def predict_image(input_image, name, age, medical_record, sex):
     if input_image is None:
@@ -36,41 +43,40 @@ def predict_image(input_image, name, age, medical_record, sex):
     # Perform prediction
     results = model(image_np)
-    # Draw bounding boxes and white circle on the image
     image_with_boxes = image_np.copy()
     raw_predictions = []
     if results[0].boxes:
         highest_confidence_result = max(results[0].boxes, key=lambda x: x.conf.item())
         class_index = highest_confidence_result.cls.item()
         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 = 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)
-        # Draw the white circle in the center of the bounding box
-        box_width = xmax - xmin
-        box_height = ymax - ymin
-        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
         thickness = 2
         # 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)
@@ -89,13 +95,47 @@ def predict_image(input_image, name, age, medical_record, sex):
     input_image.save(uploaded_folder / image_name)
     pil_image_with_boxes.save(predicted_folder / image_name)
     return pil_image_with_boxes, raw_predictions_str
 # Function to add text and watermark
 def add_text_and_watermark(image, name, age, medical_record, sex, label):
     draw = ImageDraw.Draw(image)
-    font_size = 24
     try:
         font = ImageFont.truetype("font.ttf", size=font_size)
     except IOError:
@@ -103,48 +143,75 @@ def add_text_and_watermark(image, name, age, medical_record, sex, label):
         print("Error: cannot open resource, using default font.")
     text = f"Name: {name}, Age: {age}, Medical Record: {medical_record}, Sex: {sex}, Result: {label}"
-    text_x, text_y = 20, 40
     padding = 10
     # Draw a filled rectangle for the background
     draw.rectangle(
-        [text_x - padding, text_y - padding, text_x + 500, text_y + 30 + padding],
         fill="black"
     )
     # Draw text on top of the rectangle
     draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=font)
-    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>
-            {''.join(patient_data)}
-        </body>
-    </html>
-    """
-    html_file_path = os.path.join(tempfile.gettempdir(), 'patient_info.html')
-    with open(html_file_path, 'w') as f:
-        f.write(html_content)
-    return html_file_path
 # Function to download the folders
-def download_folder(folder_path):
-    zip_path = os.path.join(tempfile.gettempdir(), f"{Path(folder_path).name}.zip")
-    shutil.make_archive(zip_path.replace('.zip', ''), 'zip', folder_path)
     return zip_path
 # Gradio Interface
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("# Cataract Detection System")
@@ -170,18 +237,16 @@ with gr.Blocks() as demo:
         download_uploaded_btn = gr.Button("Download Uploaded Images")
         download_predicted_btn = gr.Button("Download 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")
-    # Use gr.State to hold folder paths
-    uploaded_folder_state = gr.State(str(uploaded_folder))
-    predicted_folder_state = gr.State(str(predicted_folder))
-    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_folder, inputs=[uploaded_folder_state], outputs=uploaded_folder_file)
-    download_predicted_btn.click(fn=download_folder, inputs=[predicted_folder_state], outputs=predicted_folder_file)
 # Launch Gradio app
 demo.launch()

 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
 # Load YOLOv8 model
 model = YOLO("best.pt")
 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 database file if not present
+if not html_db_file.exists():
+    with open(html_db_file, 'w') as f:
+        f.write("<html><body><h1>Patient Prediction Database</h1>")
 def predict_image(input_image, name, age, medical_record, sex):
     if input_image is None:
     # Perform prediction
     results = model(image_np)
+    # Draw bounding boxes on the image
     image_with_boxes = image_np.copy()
     raw_predictions = []
     if results[0].boxes:
+        # 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)
     input_image.save(uploaded_folder / image_name)
     pil_image_with_boxes.save(predicted_folder / image_name)
+    # Convert the 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 the prediction to the HTML database
+    append_patient_info_to_html(name, age, medical_record, sex, label, predicted_image_base64)
     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)
+    # Load a larger font (adjust the size as needed)
+    font_size = 24  # Example font size
     try:
         font = ImageFont.truetype("font.ttf", size=font_size)
     except IOError:
         print("Error: cannot open resource, using default font.")
     text = f"Name: {name}, Age: {age}, Medical Record: {medical_record}, Sex: {sex}, 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 append patient info and image to HTML database
+def append_patient_info_to_html(name, age, medical_record, sex, result, predicted_image_base64):
+    html_entry = f"""
+    <div>
+        <h2>Patient Information</h2>
+        <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>
+        <p><strong>Predicted Image:</strong></p>
+        <img src="data:image/png;base64,{predicted_image_base64}" alt="Predicted Image" width="300">
+    </div>
+    <hr>
+    """
+    with open(html_db_file, 'a') as f:
+        f.write(html_entry)
 # Function to download the folders
+def download_folder(folder):
+    zip_path = os.path.join(tempfile.gettempdir(), f"{folder}.zip")
+    # Zip the folder
+    shutil.make_archive(zip_path.replace('.zip', ''), 'zip', 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)
+    # Return the current state of the HTML file with all predictions
+    return output_image, raw_result, str(html_db_file)
+# Download Functions
+def download_predicted_folder():
+    return download_folder(predicted_folder)
+def download_uploaded_folder():
+    return download_folder(uploaded_folder)
+# Launch Gradio Interface
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("# Cataract Detection System")
         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=interface, 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)
 # Launch Gradio app
 demo.launch()