Spaces:

WellnessNexus
/

ThesisCataractLatest

Sleeping

App Files Files Community

ariankhalfani commited on Sep 14

Commit

3c66f88

•

1 Parent(s): e04afcf

Create app0.py

Browse files

Files changed (1) hide show

app0.py +265 -0

app0.py ADDED Viewed

	@@ -0,0 +1,265 @@

+import gradio as gr
+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
+# Load YOLOv8 model
+model = YOLO("best.pt")
+# Create directories if not present
+uploaded_folder = Path('Uploaded_Picture')
+predicted_folder = Path('Predicted_Picture')
+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>
+        <head><title>Patient Prediction Database</title></head>
+        <body>
+        <h1>Patient Prediction Database</h1>
+        <table border="1" style="width:100%; border-collapse: collapse; text-align: center;">
+        <thead>
+            <tr>
+                <th>Name</th>
+                <th>Age</th>
+                <th>Medical Record</th>
+                <th>Sex</th>
+                <th>Result</th>
+                <th>Predicted Image</th>
+            </tr>
+        </thead>
+        <tbody>
+        """)
+def predict_image(input_image, name, age, medical_record, sex):
+    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 = []
+    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)
+        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, age, medical_record, sex, label)
+    # Save images to directories
+    image_name = f"{name}-{age}-{sex}-{medical_record}.png"
+    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:
+        font = ImageFont.load_default()
+        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
+def append_patient_info_to_html(name, age, medical_record, sex, result, predicted_image_base64):
+    # Add the patient info to the HTML table
+    html_entry = f"""
+    <tr>
+        <td>{name}</td>
+        <td>{age}</td>
+        <td>{medical_record}</td>
+        <td>{sex}</td>
+        <td>{result}</td>
+        <td><img src="data:image/png;base64,{predicted_image_base64}" alt="Predicted Image" width="150"></td>
+    </tr>
+    """
+    with open(html_db_file, 'a') as f:
+        f.write(html_entry)
+    # Ensure we only add the closing tags once
+    if "</tbody></table></body></html>" not in open(html_db_file).read():
+        with open(html_db_file, 'a') as f:
+            f.write("""
+            </tbody>
+            </table>
+            </body>
+            </html>
+            """)
+    return str(html_db_file)  # Return the HTML file path for download
+# 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 Blocks Interface
+def demo():
+    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.")
+            gr.Markdown("This application uses YOLOv8 with mAP=0.981")
+        with gr.Column():
+            name = gr.Textbox(label="Name")
+            age = gr.Number(label="Age")
+            medical_record = gr.Textbox(label="Medical Record")
+            sex = gr.Radio(["Male", "Female"], label="Sex")
+            input_image = gr.Image(type="pil", label="Upload an Image", image_mode="RGB")
+        with gr.Column():
+            submit_btn = gr.Button("Submit")
+            output_image = gr.Image(type="pil", label="Predicted Image")
+        with gr.Row():
+            raw_result = gr.Textbox(label="Prediction Result")
+        with gr.Row():
+            download_html_btn = gr.Button("Download Patient Information (HTML)")
+            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=lambda name, age, medical_record, sex, raw_result: html_db_file,
+                                inputs=[name, age, medical_record, sex, raw_result], outputs=patient_info_file)
+        download_uploaded_btn.click(fn=lambda: download_folder('Uploaded_Picture'), outputs=uploaded_folder_file)
+        download_predicted_btn.click(fn=lambda: download_folder('Predicted_Picture'), outputs=predicted_folder_file)
+    # Launch Gradio app
+    demo.launch()
+# Calling the demo function to initiate the interface
+demo()