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()