it22233530/reefsense

1from ultralytics import YOLO2import numpy as np3import gradio as gr4import cv25 6# Load models7coral_model = YOLO("coral_best.pt")8bleach_model = YOLO("coralbleaching_yolo11s_seg_best.pt")9 10def detect_bleaching_color(coral_crop):11    """12    Detect bleaching based on color (whiteness)13    """14 15    # Convert to HSV16    hsv = cv2.cvtColor(coral_crop, cv2.COLOR_BGR2HSV)17 18    # Define white / pale coral range19    lower_white = np.array([0, 0, 180])20    upper_white = np.array([180, 50, 255])21 22    white_mask = cv2.inRange(hsv, lower_white, upper_white)23 24    # Calculate percentage of white pixels25    white_pixels = np.sum(white_mask > 0)26    total_pixels = coral_crop.shape[0] * coral_crop.shape[1]27 28    white_ratio = white_pixels / total_pixels29 30    return white_ratio31 32 33def predict(image):34 35    original_image = image.copy()36 37    coral_results = coral_model(image)38    coral_masks = coral_results[0].masks39 40    total_coral = 041    total_bleached = 042 43    if coral_masks is None:44        return original_image, {45            "coral_detected": 0,46            "bleaching_detected": 0,47            "bleaching_percentage": 048        }49 50    for mask in coral_masks.data:51 52        total_coral += 153 54        mask_np = mask.cpu().numpy()55 56        mask_resized = cv2.resize(57            mask_np,58            (image.shape[1], image.shape[0])59        )60 61        binary_mask = mask_resized > 0.562 63        # Extract coral region only64        coral_crop = image.copy()65        coral_crop[~binary_mask] = 066 67        # 🔍 Color-based bleaching detection68        white_ratio = detect_bleaching_color(coral_crop)69 70        # 🤖 Model-based detection71        bleach_results = bleach_model(coral_crop)72        model_detected = len(bleach_results[0].boxes) > 073 74        # 🧠 Combined decision75        is_bleached = (white_ratio > 0.25) or model_detected76 77        # 🎯 Severity classification78        if white_ratio > 0.6:79            severity = "Severe"80        elif white_ratio > 0.3:81            severity = "Moderate"82        elif white_ratio > 0.15:83            severity = "Mild"84        else:85            severity = "Healthy"86 87        if is_bleached:88            total_bleached += 189            color = (255, 0, 0)  # RED90        else:91            color = (0, 255, 0)  # GREEN92 93        # Create colored mask overlay94        colored_mask = np.zeros_like(image)95        colored_mask[binary_mask] = color96 97        original_image = cv2.addWeighted(98            original_image, 1,99            colored_mask, 0.4,100            0101        )102 103        # Optional: draw severity text104        y, x = np.where(binary_mask)105        if len(x) > 0 and len(y) > 0:106            cx, cy = int(np.mean(x)), int(np.mean(y))107            cv2.putText(original_image, severity,108                        (cx, cy),109                        cv2.FONT_HERSHEY_SIMPLEX,110                        0.5, color, 2)111 112    bleaching_percentage = round(113        (total_bleached / total_coral) * 100, 2114    ) if total_coral > 0 else 0115 116    return original_image, {117        "coral_detected": total_coral,118        "bleaching_detected": total_bleached,119        "bleaching_percentage": bleaching_percentage120    }121 122 123iface = gr.Interface(124    fn=predict,125    inputs=gr.Image(type="numpy"),126    outputs=[127        gr.Image(type="numpy", label="Detection Result"),128        gr.JSON(label="Statistics")129    ]130)131 132iface.launch()