glaucoma.py

import numpy as np
from PIL import Image, ImageDraw, ImageFont
import cv2
from ultralytics import YOLO
from database import save_prediction_to_db

# Load YOLO models
try:
    yolo_model_glaucoma = YOLO('best-glaucoma-seg.pt')
    yolo_model_od = YOLO("best-glaucoma-od.pt")
    print("YOLO models loaded successfully.")
except Exception as e:
    print(f"Error loading YOLO models: {e}")

def calculate_area(mask):
    area = np.sum(mask > 0.5)
    print(f"Calculated area: {area}")
    return area

def classify_ddls(rim_to_disc_ratio):
    if rim_to_disc_ratio >= 0.5:
        stage = 0  # Non Glaucomatous
    elif 0.4 <= rim_to_disc_ratio < 0.5:
        stage = 1
    elif 0.3 <= rim_to_disc_ratio < 0.4:
        stage = 2
    elif 0.2 <= rim_to_disc_ratio < 0.3:
        stage = 3
    elif 0.1 <= rim_to_disc_ratio < 0.2:
        stage = 4
    elif 0.0 < rim_to_disc_ratio < 0.1:
        stage = 5
    else:
        stage = 6
    print(f"Classified DDLS stage: {stage}")
    return stage

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

def predict_and_visualize_glaucoma(image, mask_threshold=0.5):
    try:
        pil_image = Image.fromarray(image)
        orig_size = pil_image.size
        results = yolo_model_glaucoma(pil_image)

        raw_response = str(results)
        print(f"YOLO results: {raw_response}")
        masked_image = np.array(pil_image)
        mask_image = np.zeros_like(masked_image)

        cup_mask, disk_mask = None, None

        if len(results) > 0:
            result = results[0]
            if hasattr(result, 'masks') and result.masks is not None and len(result.masks) > 0:
                for mask_data in result.masks.data:
                    mask = np.array(mask_data.cpu().squeeze().numpy())
                    mask_resized = cv2.resize(mask, orig_size, interpolation=cv2.INTER_NEAREST)

                    if np.sum(mask_resized) > np.sum(disk_mask if disk_mask is not None else 0):
                        cup_mask = disk_mask
                        disk_mask = mask_resized
                    else:
                        cup_mask = mask_resized

        if cup_mask is not None and disk_mask is not None:
            area_cup = calculate_area(cup_mask)
            area_disk = calculate_area(disk_mask)
            rim_area = area_disk - area_cup
            print(f"Area cup: {area_cup}, Area disk: {area_disk}, Rim area: {rim_area}")

            rim_to_disc_ratio = rim_area / area_disk if area_disk > 0 else 0
            print(f"Rim to disc ratio: {rim_to_disc_ratio}")
            ddls_stage = classify_ddls(rim_to_disc_ratio)

            combined_image = np.array(pil_image)

            # Create RGBA version of the original image
            combined_image_rgba = cv2.cvtColor(combined_image, cv2.COLOR_RGB2RGBA)

            # Create transparent masks
            cup_mask_rgba = np.zeros_like(combined_image_rgba)
            cup_mask_rgba[:, :, 0] = 0    # Red channel
            cup_mask_rgba[:, :, 1] = 0    # Green channel
            cup_mask_rgba[:, :, 2] = 255  # Blue channel
            cup_mask_rgba[:, :, 3] = 128  # Alpha channel (50% transparency)

            disk_mask_rgba = np.zeros_like(combined_image_rgba)
            disk_mask_rgba[:, :, 0] = 255  # Red channel
            disk_mask_rgba[:, :, 1] = 0    # Green channel
            disk_mask_rgba[:, :, 2] = 0    # Blue channel
            disk_mask_rgba[:, :, 3] = 128  # Alpha channel (50% transparency)

            # Apply masks to the original image with transparency
            cup_mask_indices = cup_mask > mask_threshold
            disk_mask_indices = disk_mask > mask_threshold

            combined_image_rgba[cup_mask_indices] = (0.5 * combined_image_rgba[cup_mask_indices] + 0.5 * cup_mask_rgba[cup_mask_indices]).astype(np.uint8)
            combined_image_rgba[disk_mask_indices] = (0.5 * combined_image_rgba[disk_mask_indices] + 0.5 * disk_mask_rgba[disk_mask_indices]).astype(np.uint8)

            # Convert to PIL image for drawing
            combined_pil_image = Image.fromarray(combined_image_rgba)

            # Add text to the image
            draw = ImageDraw.Draw(combined_pil_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"Area cup: {area_cup}\nArea disk: {area_disk}\nRim area: {rim_area}\nRim to disc ratio: {rim_to_disc_ratio:.2f}\nDDLS stage: {ddls_stage}"
            text_x = 20
            text_y = 40

            draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=font)

            # Add watermark
            combined_pil_image = add_watermark(combined_pil_image)

            return np.array(combined_pil_image), area_cup, area_disk, rim_area, rim_to_disc_ratio, ddls_stage

        print("No detected regions")
        return np.zeros_like(image), 0, 0, 0, 0, "No detected regions"
    except Exception as e:
        print("Error:", e)
        return np.zeros_like(image), 0, 0, 0, 0, str(e)

def combined_prediction_glaucoma(image, mask_threshold):
    segmented_image, cup_area, disk_area, rim_area, rim_to_disc_ratio, ddls_stage = predict_and_visualize_glaucoma(image, mask_threshold)
    print(f"Segmented image: {segmented_image.shape}")
    print(f"Cup area: {cup_area}, Disk area: {disk_area}, Rim area: {rim_area}")
    print(f"Rim to disc ratio: {rim_to_disc_ratio}, DDLS stage: {ddls_stage}")

    return segmented_image, cup_area, disk_area, rim_area, rim_to_disc_ratio, ddls_stage

def submit_to_db(image, cup_area, disk_area, rim_area, rim_to_disc_ratio, ddls_stage):
    try:
        # Convert the image from numpy array to PIL image
        pil_image = Image.fromarray(np.uint8(image))
        save_prediction_to_db(pil_image, cup_area, disk_area, rim_area, rim_to_disc_ratio, ddls_stage)
        return "Values successfully saved to database.", ""
    except Exception as e:
        print(f"Error saving to database: {e}")
        return f"Error saving to database: {e}", ""

def predict_image(input_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 = yolo_model_od(image_np)

    # Draw bounding boxes on the image
    image_with_boxes = image_np.copy()
    raw_predictions = []
    for result in results[0].boxes:
        confidence = result.conf.item()  # Convert tensor to standard Python type
        label = "Glaucoma" if confidence > 0.5 else "Normal"  # Set label based on confidence
        
        xmin, ymin, xmax, ymax = map(int, result.xyxy[0])
        
        # Draw black rectangle as background for text
        text = f'{label} {confidence:.2f}'
        font_scale = 1.0  # Increased font scale
        font_thickness = 2  # Increased font thickness
        (w, h), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
        cv2.rectangle(image_with_boxes, (xmin, ymin - h - baseline), (xmin + w, ymin), (0, 0, 0), -1)
        
        cv2.putText(image_with_boxes, text, (xmin, ymin - baseline), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness)

        # Draw thicker bounding box
        box_thickness = 3  # Increased box thickness
        cv2.rectangle(image_with_boxes, (xmin, ymin), (xmax, ymax), (0, 255, 0), box_thickness)

        raw_predictions.append(f"Label: {label}, Confidence: {confidence:.2f}, Box: [{xmin}, {ymin}, {xmax}, {ymax}]")

    raw_predictions_str = "\n".join(raw_predictions)
    
    # Add watermark to the final image with boxes
    pil_image_with_boxes = Image.fromarray(image_with_boxes)
    pil_image_with_boxes = add_watermark(pil_image_with_boxes)
    image_with_boxes = np.array(pil_image_with_boxes)
    
    return image_with_boxes, raw_predictions_str