cataract.py

import os
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
import sqlite3
from io import BytesIO
from scipy.stats import norm

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

def calculate_ratios(red_values, green_values, blue_values, total_pixels):
    if total_pixels == 0:
        return 0, 0, 0

    red_ratio = np.sum(red_values) / total_pixels
    green_ratio = np.sum(green_values) / total_pixels
    blue_ratio = np.sum(blue_values) / total_pixels

    total_ratio = red_ratio + green_ratio + blue_ratio

    if total_ratio > 0:
        red_quantity = (red_ratio / total_ratio) * 255
        green_quantity = (green_ratio / total_ratio) * 255
        blue_quantity = (blue_ratio / total_ratio) * 255
    else:
        red_quantity, green_quantity, blue_quantity = 0, 0, 0

    return red_quantity, green_quantity, blue_quantity

def cataract_staging(red_quantity, green_quantity, blue_quantity):
    # Assuming you have already defined your mean and std for each class and each RGB channel
    # Example mean and std based on earlier discussion
    mean_mature_red = 73.37
    std_mature_red = (90.12 - 41.49) / 4
    mean_mature_green = 89.48
    std_mature_green = (97.67 - 83.39) / 4
    mean_mature_blue = 92.15
    std_mature_blue = (117.82 - 75.37) / 4
    
    mean_normal_red = 67.84
    std_normal_red = (107.02 - 56.19) / 4
    mean_normal_green = 84.85
    std_normal_green = (89.89 - 80.74) / 4
    mean_normal_blue = 102.31
    std_normal_blue = (111.34 - 65.58) / 4
    
    mean_immature_red = 68.83
    std_immature_red = (85.95 - 41.49) / 4
    mean_immature_green = 89.43
    std_immature_green = (97.67 - 83.39) / 4
    mean_immature_blue = 96.74
    std_immature_blue = (117.82 - 78.41) / 4

    # Calculate likelihoods for each class
    likelihood_mature = (
        norm.pdf(red_quantity, mean_mature_red, std_mature_red) *
        norm.pdf(green_quantity, mean_mature_green, std_mature_green) *
        norm.pdf(blue_quantity, mean_mature_blue, std_mature_blue)
    )
    
    likelihood_normal = (
        norm.pdf(red_quantity, mean_normal_red, std_normal_red) *
        norm.pdf(green_quantity, mean_normal_green, std_normal_green) *
        norm.pdf(blue_quantity, mean_normal_blue, std_normal_blue)
    )
    
    likelihood_immature = (
        norm.pdf(red_quantity, mean_immature_red, std_immature_red) *
        norm.pdf(green_quantity, mean_immature_green, std_immature_green) *
        norm.pdf(blue_quantity, mean_immature_blue, std_immature_blue)
    )

    # Define prior probabilities (assuming equal prior for simplicity)
    prior_mature = 1/3
    prior_normal = 1/3
    prior_immature = 1/3

    # Apply Bayes' theorem to compute posterior probabilities
    posterior_mature = likelihood_mature * prior_mature
    posterior_normal = likelihood_normal * prior_normal
    posterior_immature = likelihood_immature * prior_immature

    # Determine the stage based on maximum posterior probability
    stages = {
        posterior_mature: "Mature",
        posterior_normal: "Normal",
        posterior_immature: "Immature"
    }
    max_posterior = max(posterior_mature, posterior_normal, posterior_immature)
    stage = stages[max_posterior]

    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(image):
    try:
        pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
        orig_size = pil_image.size
        results = yolo_model_cataract(pil_image)

        raw_response = str(results)
        masked_image = np.array(pil_image)
        mask_image = np.zeros_like(masked_image)

        red_quantity, green_quantity, blue_quantity = 0, 0, 0
        total_pixels = 0

        if len(results) > 0:
            result = results[0]
            if hasattr(result, 'masks') and result.masks is not None and len(result.masks) > 0:
                mask = np.array(result.masks.data.cpu().squeeze().numpy())
                mask_resized = np.array(Image.fromarray(mask).resize(orig_size, Image.NEAREST))

                red_mask = np.zeros_like(masked_image)
                red_mask[mask_resized > 0.5] = [255, 0, 0]
                alpha = 0.5
                blended_image = cv2.addWeighted(masked_image, 1 - alpha, red_mask, alpha, 0)

                pupil_pixels = np.array(pil_image)[mask_resized > 0.5]
                total_pixels = pupil_pixels.shape[0]

                red_values = pupil_pixels[:, 0]
                green_values = pupil_pixels[:, 1]
                blue_values = pupil_pixels[:, 2]

                red_quantity, green_quantity, blue_quantity = calculate_ratios(red_values, green_values, blue_values, total_pixels)
                stage = cataract_staging(red_quantity, green_quantity, blue_quantity)

                # Add text to the blended image
                combined_pil_image = Image.fromarray(blended_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"Red quantity: {red_quantity:.2f}\nGreen quantity: {green_quantity:.2f}\nBlue quantity: {blue_quantity:.2f}\nStage: {stage}"
                
                # 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
                combined_pil_image_with_watermark = add_watermark(combined_pil_image)

                return np.array(combined_pil_image_with_watermark), red_quantity, green_quantity, blue_quantity, raw_response, stage

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

def check_duplicate_entry(conn, red_quantity, green_quantity, blue_quantity, stage):
    cursor = conn.cursor()
    query = '''SELECT COUNT(*) FROM cataract_results WHERE red_quantity=? AND green_quantity=? AND blue_quantity=? AND stage=?'''
    cursor.execute(query, (red_quantity, green_quantity, blue_quantity, stage))
    count = cursor.fetchone()[0]
    return count > 0

def save_cataract_prediction_to_db(image, red_quantity, green_quantity, blue_quantity, stage):
    database = "cataract_results.db"
    conn = create_connection(database)
    if conn:
        create_cataract_table(conn)
        
        # Check for duplicate entries
        if check_duplicate_entry(conn, red_quantity, green_quantity, blue_quantity, stage):
            conn.close()
            return "Duplicate entry found, not saving.", "Duplicate entry detected."
        
        sql = '''INSERT INTO cataract_results(image, red_quantity, green_quantity, blue_quantity, stage) VALUES(?,?,?,?,?)'''
        cur = conn.cursor()
        
        # Convert the image to bytes
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_bytes = buffered.getvalue()
        
        cur.execute(sql, (img_bytes, red_quantity, green_quantity, blue_quantity, stage))
        conn.commit()
        conn.close()
        return "Data saved successfully", f"Red: {red_quantity}, Green: {green_quantity}, Blue: {blue_quantity}, Stage: {stage}"

    return "Failed to save data", "No connection to the database."

def combined_prediction(image):
    blended_image, red_quantity, green_quantity, blue_quantity, raw_response, stage = predict_and_visualize(image)
    save_message, debug_info = save_cataract_prediction_to_db(Image.fromarray(blended_image), red_quantity, green_quantity, blue_quantity, stage)
    return blended_image, red_quantity, green_quantity, blue_quantity, raw_response, stage, save_message, debug_info

def create_connection(db_file):
    """ Create a database connection to the SQLite database """
    conn = None
    try:
        conn = sqlite3.connect(db_file)
        return conn
    except sqlite3.Error as e:
        print(e)
    return conn

def create_cataract_table(conn):
    """ Create the cataract results table if it does not exist """
    create_table_sql = """ CREATE TABLE IF NOT EXISTS cataract_results (
                            id integer PRIMARY KEY,
                            image blob,
                            red_quantity real,
                            green_quantity real,
                            blue_quantity real,
                            stage text
                        ); """
    try:
        cursor = conn.cursor()
        cursor.execute(create_table_sql)
    except sqlite3.Error as e:
        print(e)

def predict_object_detection(image):
    try:
        image_np = np.array(image)
        results = yolo_model_object_detection(image_np)

        image_with_boxes = image_np.copy()
        raw_predictions = []
        for result in results[0].boxes:
            label = "Normal" if result.cls.item() == 1 else "Cataract"
            confidence = result.conf.item()
            xmin, ymin, xmax, ymax = map(int, result.xyxy[0])
            cv2.rectangle(image_with_boxes, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
            
            font_scale = 1.0
            thickness = 2
            text = f'{label} {confidence:.2f}'
            (text_width, text_height), baseline = cv2.getTextSize(text, 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)
            cv2.putText(image_with_boxes, text, (xmin, ymin - baseline), 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 image_with_boxes to PIL image and add watermark
        image_with_boxes_pil = Image.fromarray(image_with_boxes)
        image_with_boxes_pil_with_watermark = add_watermark(image_with_boxes_pil)

        return np.array(image_with_boxes_pil_with_watermark), raw_predictions_str
    except Exception as e:
        print("Error in object detection:", e)
        return np.zeros_like(image), str(e)