Spaces:

WellnessNexus
/

ThesisCataractLatest

Sleeping

File size: 9,887 Bytes

3c66f88

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()