app.py

import io
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 transformers import Mask2FormerImageProcessor, Mask2FormerForUniversalSegmentation

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

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

# NOTE: Ensure the folder "./mask2former-manual-save" is actually uploaded to your HF Space!
try:
    model = Mask2FormerForUniversalSegmentation.from_pretrained(
        "./mask2former-manual-save",
        ignore_mismatched_sizes=True
    )
except Exception as e:
    print(f"Warning: Local model folder not found, falling back to base model. Error: {e}")
    model = Mask2FormerForUniversalSegmentation.from_pretrained(
        MODEL_ID,
        ignore_mismatched_sizes=True
    )

SHOW_CATEGORIES = {"Spore": True, "Conidal": True, "Hypha": True, "Suspected": False}
LABEL_MAPPING = {
    "spore_asp": "spore",
    "hypha_asp": "hyphae",
    "conidalhead_asp": "conidial head",
    "suspected_asp": "hyphae"
}
CLASS_COLORS = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1]]

# Sliding Window Settings
WINDOW_SIZE = 512
STRIDE = 384

# ==========================================
# 2. AUTO-SEGMENTATION INFERENCE FUNCTION
# ==========================================
# INCREASED duration: Sliding window over large images easily exceeds 7s. 
@spaces.GPU(duration=7) 
def run_mask2former_inference(images_state, conf_threshold, progress=gr.Progress()):
    """Runs sliding window segmentation inference on the full-res image using ZeroGPU."""
    if not images_state:
        return None, [], [], "No image available to run segmentation.", "0.00%"

    progress(0.1, desc="Loading image onto GPU...")
    
    # REQUIREMENT 1 FIX: Strictly use the original unmodified image for inference
    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()

    id2label = model.config.id2label
    global_masks_per_class = {}
    
    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 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")
            pixel_values = inputs["pixel_values"].to(device)
            
            with torch.no_grad():
                outputs = model(pixel_values=pixel_values)
                
            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()
                class_id = confidences.argmax().item()
                label = id2label.get(class_id, str(class_id))
                
                # REQUIREMENT 2 FIX: Use dynamic conf_threshold from slider
                if max_conf >= conf_threshold and SHOW_CATEGORIES.get(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
                    )
                    mask = m_resized.sigmoid().squeeze().numpy()
                    mask_binary_local = mask > 0.8

                    if np.any(mask_binary_local):
                        if class_id not in global_masks_per_class:
                            global_masks_per_class[class_id] = {
                                "max_conf": max_conf,
                                "mask": np.zeros((target_h, target_w), dtype=bool)
                            }
                        else:
                            global_masks_per_class[class_id]["max_conf"] = max(
                                global_masks_per_class[class_id]["max_conf"], max_conf
                            )
                        global_masks_per_class[class_id]["mask"][y_min:y_max, x_min:x_max] |= mask_binary_local

    progress(0.85, desc="Plotting contours and stitched boundaries...")
    
    # Render Plot via Matplotlib (Safe Headless Mode)
    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 = [
        (cid, data["max_conf"], data["mask"]) 
        for cid, data in global_masks_per_class.items() if np.any(data["mask"])
    ]
    
    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):
        label = id2label.get(class_id, str(class_id))
        color = CLASS_COLORS[class_id % len(CLASS_COLORS)]
        display_name = LABEL_MAPPING.get(label, label)
        
        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
        class_pixel_count = np.sum(mask_binary)
        class_coverage_pct = (class_pixel_count / total_pixels) * 100
        coverage_stats.append(f"• {display_name.capitalize()}: {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') # Safely clears memory for HF Spaces
    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 = np.sum(combined_total_mask)
    total_coverage_pct = (total_covered_pixels / total_pixels) * 100
    
    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!")
    
    # REQUIREMENT 1 FIX: Store results in display variables, keeping the original untouched
    images_state[0]["disp_low"] = low_segmented
    images_state[0]["disp_full"] = segmented_pil

    return low_segmented, images_state, [], "Mask2Former auto-segmentation completed.", coverage_report

# ==========================================
# 3. GRADIO ANNOTATION HELPER FUNCTIONS
# ==========================================
def load_and_crop_images(files, progress=gr.Progress()):
    """Processes and automatically crops uploaded images to a 3:2 aspect ratio."""
    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()
            
        # Refactored state formatting to separate base image and display canvas
        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 ({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 image...")
    
    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})."

# ==========================================
# 4. GRADIO INTERFACE LAYOUT
# ==========================================
with gr.Blocks() as demo:
    gr.Markdown("# Interactive Image Segmentation & Annotation Workspace")
    annotations_state = gr.State([])
    images_state = gr.State([])
    
    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="secondary")
        # REQUIREMENT 2: Added Confidence slider
        conf_threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.75, step=0.05, label="Confidence Threshold")
        rerender_btn = gr.Button("Render Manual Annotations in Full Scale", variant="primary")
        
    with gr.Row():
        with gr.Column():
            image_upload = gr.Files(
                file_types=["image"],
                file_count="multiple",
                label="Upload Images",
            )
            upload_status = gr.Textbox(label="Status")
            
        with gr.Column():
            label_dropdown = gr.Dropdown(
                choices=["Spore", "Spores", "Hyphae", "Conidial Head"],
                value="Spores",
                label="Select Label",
                interactive=True,
            )
            with gr.Row():
                new_label_input = gr.Textbox(label="Add New Label", scale=2)
                add_label_btn = gr.Button("Add to Dropdown", scale=1)
                
            label_placement = gr.Radio(
                choices=["Top Left", "Top Right", "Bottom Left", "Bottom Right"],
                value="Bottom Right",
                label="Label Box Position (Relative to Click Target)",
                interactive=True,
            )
            
    with gr.Row():
        click_log = gr.Textbox(label="Action Log", scale=2)
        coverage_box = gr.Textbox(label="Model Area Coverage (%)", scale=1, interactive=False)

    # --- Event Binding ---
    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],
        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()