app.py

import io
import json
import os
import cv2
import gradio as gr
import numpy as np
import matplotlib
matplotlib.use('Agg')  # CRITICAL: Makes matplotlib thread-safe for Hugging Face web servers
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import spaces
from PIL import Image
from pycocotools import mask as mask_utils
from transformers import Mask2FormerImageProcessor, Mask2FormerForUniversalSegmentation

# ==========================================
# 1. MODEL INITIALIZATION & GLOBAL CONFIG
# ==========================================
MODEL_ID = "facebook/mask2former-swin-tiny-coco-instance"

# Load processor globally
processor = Mask2FormerImageProcessor.from_pretrained(MODEL_ID)

# Attempt to load local trained model checkpoints with fallbacks
model = None
checkpoint_paths = ["./mask2former2", "./mask2former-manual-save", MODEL_ID]

for path in checkpoint_paths:
    try:
        model = Mask2FormerForUniversalSegmentation.from_pretrained(
            path,
            ignore_mismatched_sizes=True
        )
        print(f"Successfully loaded model checkpoint from: '{path}'")
        break
    except Exception as e:
        print(f"Info: Could not load model from '{path}'. Reason: {e}")

if model is None:
    raise RuntimeError("Failed to load any valid Mask2Former model checkpoint.")

# Model output integer class mapping (0-indexed)
ID_TO_LABEL = {
    0: "Conidial Head",
    1: "Hyphae",
    2: "Free Spore",
    3: "Spore Clump",
    4: "Debris",
    5: "Fiber",
    6: "Bubble",
}

SHOW_CATEGORIES = {label: True for label in ID_TO_LABEL.values()}

# Distinct color palette per class [R, G, B] normalized to 0.0-1.0
CLASS_COLORS = [
    [1.0, 0.0, 0.0],  # 0: Conidial Head -> Red
    [0.0, 1.0, 0.0],  # 1: Hyphae        -> Green
    [0.0, 0.0, 1.0],  # 2: Free Spore    -> Blue
    [1.0, 1.0, 0.0],  # 3: Spore Clump   -> Yellow
    [1.0, 0.0, 1.0],  # 4: Debris        -> Magenta
    [0.0, 1.0, 1.0],  # 5: Fiber         -> Cyan
    [0.5, 0.5, 0.5],  # 6: Bubble        -> Gray
]

# Sliding Window Hyperparameters
WINDOW_SIZE = 512
STRIDE = 384


# ==========================================
# 2. HELPER UTILITIES
# ==========================================
def binary_mask_to_rle(binary_mask: np.ndarray) -> dict:
    """Converts a 2D boolean NumPy mask [H, W] into standard COCO RLE format."""
    fortran_mask = np.asfortranarray(binary_mask.astype(np.uint8))
    rle = mask_utils.encode(fortran_mask)
    rle["counts"] = rle["counts"].decode("utf-8")  # Decode bytes to UTF-8 for JSON
    return rle


# ==========================================
# 3. INFERENCE ENGINE (WITH PROBABILITY ACCUMULATION & RLE)
# ==========================================
@spaces.GPU(duration=30)
def run_mask2former_inference(images_state, conf_threshold, progress=gr.Progress()):
    """Runs sliding-window inference with probability accumulation and generates a COCO RLE JSON file."""
    empty_json = {"error": "No image loaded."}
    if not images_state:
        return None, [], [], "No image available to run segmentation.", "0.00%", empty_json, None

    progress(0.1, desc="Loading image onto GPU...")

    # Load full-resolution image from state
    original_image = images_state[0]["orig_full"]
    img = np.array(original_image)
    target_h, target_w = img.shape[:2]
    total_pixels = target_h * target_w
    img_display = (img - img.min()) / (img.max() - img.min() + 1e-8)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    num_classes = len(ID_TO_LABEL)

    # Accumulation matrices for maximum probability projection
    accumulated_probs = np.zeros((num_classes, target_h, target_w), dtype=np.float32)
    class_max_confs = {cid: 0.0 for cid in range(num_classes)}

    y_steps = list(range(0, target_h, STRIDE))
    x_steps = list(range(0, target_w, STRIDE))
    total_steps = max(1, len(y_steps) * len(x_steps))
    step_counter = 0

    # Sliding Window Processing Loop
    for y_min in y_steps:
        y_max = min(y_min + WINDOW_SIZE, target_h)
        if (y_max - y_min) < 64:
            continue

        for x_min in x_steps:
            x_max = min(x_min + WINDOW_SIZE, target_w)
            if (x_max - x_min) < 64:
                continue

            step_counter += 1
            progress(
                0.1 + 0.7 * (step_counter / total_steps),
                desc=f"Segmenting patch {step_counter}/{total_steps}..."
            )

            crop_image = original_image.crop((x_min, y_min, x_max, y_max))
            crop_h, crop_w = crop_image.size[1], crop_image.size[0]

            inputs = processor(images=crop_image, return_tensors="pt").to(device)

            with torch.no_grad():
                outputs = model(**inputs)

            # Softmax class queries (exclude background channel :-1)
            probas = outputs.class_queries_logits.softmax(-1)[0, :, :-1].cpu()
            mask_logits = outputs.masks_queries_logits[0].cpu()

            for i in range(len(probas)):
                confidences = probas[i]
                max_conf = confidences.max().item()
                raw_class_id = confidences.argmax().item()

                # Align 1-based or 0-based label indices
                if raw_class_id in ID_TO_LABEL:
                    class_id = raw_class_id
                elif (raw_class_id - 1) in ID_TO_LABEL:
                    class_id = raw_class_id - 1
                else:
                    continue

                display_label = ID_TO_LABEL[class_id]

                if max_conf >= conf_threshold and SHOW_CATEGORIES.get(display_label, True):
                    m_logit = mask_logits[i].unsqueeze(0).unsqueeze(0)
                    m_resized = F.interpolate(
                        m_logit, size=(crop_h, crop_w), mode="bilinear", align_corners=False
                    )
                    prob_map = m_resized.sigmoid().squeeze().numpy() * max_conf

                    # Maximum Continuous Probability Accumulation
                    accumulated_probs[class_id, y_min:y_max, x_min:x_max] = np.maximum(
                        accumulated_probs[class_id, y_min:y_max, x_min:x_max],
                        prob_map
                    )
                    class_max_confs[class_id] = max(class_max_confs[class_id], max_conf)

    progress(0.85, desc="Generating COCO RLE masks and plotting boundaries...")

    # Data Structure for JSON Output
    json_output = {
        "image_metadata": {
            "width": int(target_w),
            "height": int(target_h),
            "total_pixels": int(total_pixels),
            "confidence_threshold_used": float(conf_threshold)
        },
        "summary": {
            "total_coverage_pct": 0.0,
            "detected_classes_count": 0
        },
        "annotations": []
    }

    # Render Visualization via Matplotlib
    fig, ax = plt.subplots(figsize=(14, 10))
    ax.imshow(img_display)
    label_x_position = target_w + (target_w * 0.03)
    ax.set_xlim(0, target_w + (target_w * 0.25))
    ax.set_ylim(target_h, 0)

    valid_predictions = []
    for class_id in range(num_classes):
        prob_map = accumulated_probs[class_id]
        mask_binary = prob_map >= 0.5  # Binarize accumulated probabilities
        if np.any(mask_binary):
            valid_predictions.append((class_id, class_max_confs[class_id], mask_binary))

    def get_top_y_coordinate(item):
        mask_binary = item[2]
        y_indices, _ = np.where(mask_binary)
        return y_indices.min() if len(y_indices) > 0 else target_h

    valid_predictions.sort(key=get_top_y_coordinate)
    start_y = target_h * 0.05
    y_spacing = target_h * 0.05

    combined_total_mask = np.zeros((target_h, target_w), dtype=bool)
    coverage_stats = []

    for idx, (class_id, max_conf, mask_binary) in enumerate(valid_predictions):
        display_name = ID_TO_LABEL.get(class_id, f"Class {class_id}")
        color = CLASS_COLORS[class_id % len(CLASS_COLORS)]

        # Calculate bounding box [xmin, ymin, width, height]
        y_indices, x_indices = np.where(mask_binary)
        bbox = [
            int(x_indices.min()),
            int(y_indices.min()),
            int(x_indices.max() - x_indices.min()),
            int(y_indices.max() - y_indices.min())
        ]

        class_pixel_count = int(np.sum(mask_binary))
        class_coverage_pct = round((class_pixel_count / total_pixels) * 100, 2)

        # Convert Binary Mask to COCO RLE
        rle_mask = binary_mask_to_rle(mask_binary)

        # Append COCO Annotation to JSON output
        json_output["annotations"].append({
            "id": idx + 1,
            "category_id": int(class_id),
            "category_name": display_name,
            "score": round(float(max_conf), 4),
            "area": class_pixel_count,
            "coverage_percentage": class_coverage_pct,
            "bbox": bbox,
            "segmentation": rle_mask
        })

        # Draw smooth contours on overlay
        ax.contour(mask_binary, levels=[0.5], colors=[color], linewidths=2.0)
        assigned_y_position = start_y + (idx * y_spacing)

        ax.text(
            label_x_position, assigned_y_position, f"{display_name} ({max_conf:.2f})",
            color='white', fontsize=10, fontweight='bold', ha='left', va='center',
            bbox=dict(facecolor=color, alpha=0.8, edgecolor='none', boxstyle='round,pad=0.5')
        )

        combined_total_mask |= mask_binary
        coverage_stats.append(f"• {display_name}: {class_coverage_pct:.2f}%")

    ax.axis('off')
    plt.title("Mask2Former Auto-Segmentation Results", fontsize=12, pad=15)
    plt.tight_layout()

    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', dpi=150)
    plt.close('all')
    buf.seek(0)

    segmented_pil = Image.open(buf).convert("RGB")
    fw, fh = segmented_pil.size
    low_w = 800
    low_h = int(fh * (800 / fw))
    low_segmented = segmented_pil.resize((low_w, low_h), Image.Resampling.LANCZOS)

    total_covered_pixels = int(np.sum(combined_total_mask))
    total_coverage_pct = round((total_covered_pixels / total_pixels) * 100, 2)

    json_output["summary"]["total_coverage_pct"] = total_coverage_pct
    json_output["summary"]["detected_classes_count"] = len(valid_predictions)

    # Write temporary file for download button
    json_filepath = "segmentation_results.json"
    with open(json_filepath, "w") as f:
        json.dump(json_output, f, indent=2)

    coverage_report = f"Total Coverage: {total_coverage_pct:.2f}%\n" + "\n".join(coverage_stats)
    if total_coverage_pct == 0:
        coverage_report = "Total Coverage: 0.00%\nNo segments detected."

    progress(1.0, desc="Segmentation complete!")

    images_state[0]["disp_low"] = low_segmented
    images_state[0]["disp_full"] = segmented_pil

    return (
        low_segmented,
        images_state,
        [],  # Clear manual annotations on fresh auto-segment
        "Mask2Former auto-segmentation completed.",
        coverage_report,
        json_output,     # Rendered in gr.JSON viewer
        json_filepath    # Output file path for gr.File / DownloadButton
    )


# ==========================================
# 4. GRADIO ANNOTATION & INTERFACE HELPERS
# ==========================================
def load_and_crop_images(files, progress=gr.Progress()):
    """Loads and crops images to a 3:2 ratio with low-res scaling for interactive responsiveness."""
    if not files:
        return None, [], [], "No images uploaded."

    processed_images = []
    target_ratio = 3 / 2
    file_list = files if isinstance(files, list) else [files]
    total_files = len(file_list)

    for idx, f in enumerate(file_list):
        progress((idx / total_files) * 0.5, desc=f"Loading image {idx + 1}/{total_files}...")

        if isinstance(f, Image.Image):
            pil_img = f.convert("RGB")
        else:
            file_path = f.name if hasattr(f, "name") else f
            pil_img = Image.open(file_path).convert("RGB")

        img = np.array(pil_img)
        h, w, _ = img.shape

        if w / h > target_ratio:
            target_w = int(h * target_ratio)
            target_h = h
        else:
            target_w = w
            target_h = int(w / target_ratio)

        start_x = (w - target_w) // 2
        start_y = (h - target_h) // 2

        cropped_img = img[start_y: start_y + target_h, start_x: start_x + target_w]
        full_img = Image.fromarray(cropped_img)

        progress(((idx + 0.5) / total_files), desc=f"Scaling image {idx + 1}/{total_files} to 800px...")
        fw, fh = full_img.size
        if fw > 800:
            low_w = 800
            low_h = int(fh * (800 / fw))
            low_img = full_img.resize((low_w, low_h), Image.Resampling.LANCZOS)
        else:
            low_img = full_img.copy()

        processed_images.append({
            "orig_low": low_img,
            "orig_full": full_img,
            "disp_low": low_img.copy(),
            "disp_full": full_img.copy()
        })

    first_low_res = processed_images[0]["disp_low"] if processed_images else None
    return first_low_res, processed_images, [], f"Loaded {len(processed_images)} image(s) cropped to 3:2."


def add_to_dropdown(new_text, current_value, dropdown_component):
    choices = getattr(dropdown_component, "choices", None)
    if not isinstance(choices, list) or not choices:
        choices = ["Spore", "Spores", "Hyphae", "Conidial Head"]
    updated_choices = list(choices)

    if new_text and new_text.strip():
        clean_text = new_text.strip()
        if clean_text not in updated_choices:
            updated_choices.append(clean_text)
        return gr.update(choices=updated_choices, value=clean_text, interactive=True), ""

    return gr.update(choices=updated_choices, value=current_value, interactive=True), ""


def draw_annotations_on_image(base_pil_img, annotations, scale_factor=1.0):
    annotated_img = np.array(base_pil_img).copy()
    height, width, _ = annotated_img.shape

    arrow_length = int(width * 0.05)
    head_size = max(8, int(width * 0.012))
    line_width = max(2, int(width * 0.003))
    font_scale = max(0.6, width * 0.0008)
    font_face = cv2.FONT_HERSHEY_SIMPLEX
    font_thickness = max(1, int(width * 0.0015))
    red_color = (255, 0, 0)

    for ann in annotations:
        x = int(ann["canvas_x"] * scale_factor)
        y = int(ann["canvas_y"] * scale_factor)
        current_label = ann["label"]
        position_mode = ann["position_mode"]

        if position_mode == "Top Left":
            arrow_back_x, arrow_back_y = x - arrow_length, y - arrow_length
            head_poly = np.array([[x, y], [x - head_size, y], [x, y - head_size]], np.int32)
        elif position_mode == "Top Right":
            arrow_back_x, arrow_back_y = x + arrow_length, y - arrow_length
            head_poly = np.array([[x, y], [x + head_size, y], [x, y - head_size]], np.int32)
        elif position_mode == "Bottom Left":
            arrow_back_x, arrow_back_y = x - arrow_length, y + arrow_length
            head_poly = np.array([[x, y], [x - head_size, y], [x, y + head_size]], np.int32)
        else:
            arrow_back_x, arrow_back_y = x + arrow_length, y + arrow_length
            head_poly = np.array([[x, y], [x + head_size, y], [x, y + head_size]], np.int32)

        cv2.line(annotated_img, (arrow_back_x, arrow_back_y), (x, y), red_color, thickness=line_width, lineType=cv2.LINE_AA)
        cv2.fillPoly(annotated_img, [head_poly], red_color)

        label_str = str(current_label)
        (text_w, text_h), _ = cv2.getTextSize(label_str, font_face, font_scale, font_thickness)
        padding = int(text_h * 0.4)

        text_x = arrow_back_x - text_w - padding if "Left" in position_mode else arrow_back_x + padding
        text_y = arrow_back_y if "Top" in position_mode else arrow_back_y + text_h + padding

        bg_rect_pt1 = (text_x - padding, text_y - text_h - padding)
        bg_rect_pt2 = (text_x + text_w + padding, text_y + padding // 2)

        cv2.rectangle(annotated_img, bg_rect_pt1, bg_rect_pt2, (240, 240, 240), -1)
        cv2.rectangle(annotated_img, bg_rect_pt1, bg_rect_pt2, (0, 0, 0), thickness=max(1, line_width // 2))
        cv2.putText(annotated_img, label_str, (text_x, text_y), font_face, font_scale, red_color, thickness=font_thickness, lineType=cv2.LINE_AA)

    return Image.fromarray(annotated_img)


def handle_image_click(evt: gr.SelectData, img, current_label, position_mode, current_annotations, images_state):
    if not images_state:
        return img, "No image loaded", current_annotations
    if current_annotations is None:
        current_annotations = []

    x, y = evt.index[0], evt.index[1]
    disp_low = images_state[0]["disp_low"]

    current_annotations.append({
        "canvas_x": x,
        "canvas_y": y,
        "label": current_label,
        "position_mode": position_mode
    })

    annotated_img = draw_annotations_on_image(disp_low, current_annotations, scale_factor=1.0)
    log_msg = f"Labeled '{current_label}' at click target ({x}, {y})."

    return annotated_img, log_msg, current_annotations


def render_full_resolution(images_state, current_annotations, progress=gr.Progress()):
    if not images_state:
        return None, "No image loaded to render."

    progress(0.2, desc="Rendering high-resolution vector overlay...")

    disp_low = images_state[0]["disp_low"]
    disp_full = images_state[0]["disp_full"]

    scale_factor = disp_full.width / disp_low.width
    annotated_full_res = draw_annotations_on_image(disp_full, current_annotations, scale_factor=scale_factor)

    progress(1.0, desc="Rendering complete!")
    return annotated_full_res, f"Rendered Full Scale Image ({disp_full.width}x{disp_full.height})."


# ==========================================
# 5. GRADIO APP LAYOUT & EVENT BINDINGS
# ==========================================
with gr.Blocks(title="Mask2Former Mold Segmentation Workspace") as demo:
    gr.Markdown("# Interactive Image Segmentation & Annotation Workspace")
    gr.Markdown("Fine-tuned Mask2Former mycology auto-segmentation with continuous probability accumulation & COCO RLE JSON exports.")

    annotations_state = gr.State([])
    images_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=2):
            annot_image = gr.Image(
                type="pil",
                format="png",
                height=533,
                show_label=False,
                interactive=True,
            )

            with gr.Row():
                segment_btn = gr.Button("Run Mask2Former Auto-Segmentation", variant="primary")
                conf_threshold = gr.Slider(
                    minimum=0.0, maximum=1.0, value=0.50, step=0.05, label="Confidence Threshold"
                )
                rerender_btn = gr.Button("Render Manual Annotations in Full Scale", variant="secondary")

        with gr.Column(scale=1):
            image_upload = gr.Files(
                file_types=["image"],
                file_count="multiple",
                label="Upload Images",
            )
            upload_status = gr.Textbox(label="Upload Status", interactive=False)

            label_dropdown = gr.Dropdown(
                choices=["Spore", "Spores", "Hyphae", "Conidial Head"],
                value="Spores",
                label="Select Manual Label",
                interactive=True,
            )
            with gr.Row():
                new_label_input = gr.Textbox(label="Add New Label", scale=2)
                add_label_btn = gr.Button("Add", scale=1)

            label_placement = gr.Radio(
                choices=["Top Left", "Top Right", "Bottom Left", "Bottom Right"],
                value="Bottom Right",
                label="Label Box Position (Relative to Target)",
                interactive=True,
            )

    with gr.Row():
        click_log = gr.Textbox(label="Action Log", scale=2, interactive=False)
        coverage_box = gr.Textbox(label="Model Area Coverage (%)", scale=1, interactive=False)

    # --- JSON Output & File Download Section ---
    gr.Markdown("---")
    gr.Markdown("### COCO RLE Segmentation Metadata Output")
    with gr.Row():
        json_viewer = gr.JSON(label="Segmentation JSON Output")
        json_download_file = gr.File(label="Download Segmentation JSON File", interactive=False)

    # --- Event Bindings ---
    annot_image.upload(
        load_and_crop_images,
        inputs=annot_image,
        outputs=[annot_image, images_state, annotations_state, upload_status],
        show_progress="full",
    )

    image_upload.change(
        load_and_crop_images,
        inputs=image_upload,
        outputs=[annot_image, images_state, annotations_state, upload_status],
        show_progress="full",
    )

    add_label_btn.click(
        add_to_dropdown,
        inputs=[new_label_input, label_dropdown, label_dropdown],
        outputs=[label_dropdown, new_label_input],
        show_progress="hidden",
    )

    annot_image.select(
        handle_image_click,
        inputs=[
            annot_image,
            label_dropdown,
            label_placement,
            annotations_state,
            images_state,
        ],
        outputs=[annot_image, click_log, annotations_state],
        show_progress="minimal",
    )

    segment_btn.click(
        run_mask2former_inference,
        inputs=[images_state, conf_threshold],
        outputs=[
            annot_image,
            images_state,
            annotations_state,
            click_log,
            coverage_box,
            json_viewer,
            json_download_file
        ],
        show_progress="full",
    )

    rerender_btn.click(
        render_full_resolution,
        inputs=[images_state, annotations_state],
        outputs=[annot_image, click_log],
        show_progress="full",
    )

if __name__ == "__main__":
    demo.launch()