Dense-Captioning/medsam-inference
0
1"""2HuggingFace Space for SAM / MedSAM Inference3API-compatible with Dense-Captioning-Toolkit backend4 5Deploy this to: https://huggingface.co/spaces/YOUR_USERNAME/medsam-inference6"""7import gradio as gr8import torch9import numpy as np10from PIL import Image11import io12import json13import base6414import os15import uuid16 17from huggingface_hub import hf_hub_download18 19# Import SAM components20from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator21 22# Initialize model23device = "cuda" if torch.cuda.is_available() else "cpu"24print(f"Using device: {device}")25 26# -----------------------------------------------------------------------------27# Model configuration28# -----------------------------------------------------------------------------29 30# 1) MedSAM (ViT-B) for interactive segmentation (points / boxes / multiple boxes)31# We assume medsam_vit_b.pth is committed in this repo (small enough for Spaces).32MEDSAM_CHECKPOINT = os.path.join(os.path.dirname(__file__), "medsam_vit_b.pth")33 34print("Loading MedSAM model (vit_b) for interactive segmentation...")35try:36 # MedSAM checkpoints are typically state_dicts; load and apply to a vit_b SAM backbone.37 state_dict = torch.load(MEDSAM_CHECKPOINT, map_location=device)38 medsam = sam_model_registry["vit_b"](checkpoint=None)39 medsam.load_state_dict(state_dict)40 medsam.to(device=device)41 medsam.eval()42 print("✓ MedSAM model (vit_b) loaded successfully")43except Exception as e:44 print(f"✗ Failed to load MedSAM model from {MEDSAM_CHECKPOINT}: {e}")45 raise46 47# SamPredictor for interactive segmentation (point/box prompts) using MedSAM48predictor = SamPredictor(medsam)49print("✓ SamPredictor (MedSAM) initialized for interactive segmentation")50 51 52# 2) SAM ViT-H for automatic mask generation and embedding (encode_image)53# We download this large checkpoint from a separate model repo using hf_hub_download.54MODEL_REPO_ID = "Aniketg6/dense-captioning-models"55MODEL_FILENAME = "sam_vit_h_4b8939.pth" # change if your filename is different56MODEL_TYPE = "vit_h" # using SAM ViT-H (general-purpose SAM)57 58print(f"Downloading SAM (vit_h) checkpoint `{MODEL_FILENAME}` from repo `{MODEL_REPO_ID}`...")59SAM_CHECKPOINT = hf_hub_download(60 repo_id=MODEL_REPO_ID,61 filename=MODEL_FILENAME,62)63print(f"✓ SAM (vit_h) checkpoint downloaded to: {SAM_CHECKPOINT}")64 65print("Loading SAM model (vit_h) for auto masks and embeddings...")66 67# Monkey-patch torch.load to use CPU mapping when needed68original_torch_load = torch.load69def patched_torch_load(f, *args, **kwargs):70 if "map_location" not in kwargs and device == "cpu":71 kwargs["map_location"] = "cpu"72 return original_torch_load(f, *args, **kwargs)73 74torch.load = patched_torch_load75 76try:77 # Ensure we always load onto CPU when no GPU is available78 torch.load = patched_torch_load79 sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_CHECKPOINT)80finally:81 torch.load = original_torch_load82 83sam.to(device=device)84sam.eval()85print("✓ SAM model (vit_h) loaded successfully")86 87# SamAutomaticMaskGenerator for automatic mask generation (SAM ViT-H)88mask_generator = SamAutomaticMaskGenerator(89 model=sam,90 points_per_side=16, # Lighter grid (16x16) for faster CPU + smaller responses91 pred_iou_thresh=0.7, # IoU threshold for filtering92 stability_score_thresh=0.7, # Stability threshold93 crop_n_layers=0, # Disable multi-scale crops to avoid IndexError94 crop_n_points_downscale_factor=2,95 min_mask_region_area=0 # Allow small masks (backend can filter if needed)96)97print("✓ SamAutomaticMaskGenerator (SAM vit-h) initialized for automatic segmentation")98 99 100# =============================================================================101# HELPER FUNCTIONS FOR EMBEDDINGS (STATELESS)102# =============================================================================103 104def set_predictor_features_from_embedding(embedding_tensor: torch.Tensor, image_shape: tuple):105 """106 Set SamPredictor's internal features using precomputed embedding107 108 Args:109 embedding_tensor: Precomputed embedding tensor [1, C, H, W]110 image_shape: Original image shape (height, width)111 """112 # SamPredictor stores features in self.features113 # We need to set it directly (this is a bit of a hack but necessary)114 predictor.features = embedding_tensor115 predictor.original_image_size = image_shape116 predictor.input_size = (1024, 1024) # SAM default input size117 predictor.is_image_set = True118 119 120# =============================================================================121# API FUNCTIONS - MATCHING BACKEND FORMAT (backend/app.py)122# =============================================================================123 124def encode_image(image, request_json):125 """126 Encode image using SAM image encoder and return embedding to the client.127 128 This is now a stateless API: it does NOT talk to Supabase. The caller129 (your backend) is responsible for storing the embedding if desired.130 131 Args:132 image: PIL Image133 request_json: JSON string with optional fields:134 {135 "image_id": "uuid-string" # Optional: image ID from your DB136 }137 138 Returns:139 JSON string:140 {141 "success": true/false,142 "image_id": "uuid-string" or null,143 "embedding_npy_base64": "...", # base64-encoded .npy of [C,H,W]144 "embedding_shape": [1, C, H, W]145 }146 """147 try:148 # Parse input (image_id is optional and just echoed back)149 data = json.loads(request_json) if request_json else {}150 image_id = data.get("image_id")151 152 # Convert PIL to numpy153 image_array = np.array(image)154 H, W = image_array.shape[:2]155 156 # Resize image to SAM's expected input size (1024x1024)157 from skimage import transform158 img_resized = transform.resize(159 image_array,160 (1024, 1024),161 order=3,162 preserve_range=True,163 anti_aliasing=True,164 ).astype(np.uint8)165 166 # Normalize image (SAM expects normalized input)167 img_norm = (img_resized - img_resized.min()) / np.clip(168 img_resized.max() - img_resized.min(), 1e-8, None169 )170 171 # Convert to tensor and add batch dimension172 tensor = (173 torch.tensor(img_norm)174 .float()175 .permute(2, 0, 1)176 .unsqueeze(0)177 .to(device)178 )179 180 # Encode image using SAM image encoder181 print(f"Encoding image (image_id={image_id}) original size: {W}x{H} -> 1024x1024")182 with torch.no_grad():183 embedding = sam.image_encoder(tensor)184 185 # Convert embedding to numpy [C, Hf, Wf]186 arr = embedding.squeeze(0).cpu().numpy().astype(np.float32)187 188 # Serialize as .npy in memory and base64-encode it189 buf = io.BytesIO()190 np.save(buf, arr)191 buf.seek(0)192 embedding_b64 = base64.b64encode(buf.read()).decode("utf-8")193 194 return json.dumps(195 {196 "success": True,197 "image_id": image_id,198 "embedding_npy_base64": embedding_b64,199 "embedding_shape": list(embedding.shape),200 }201 )202 except Exception as e:203 import traceback204 return json.dumps(205 {206 "success": False,207 "error": str(e),208 "traceback": traceback.format_exc(),209 }210 )211 212 213def segment_points(image, request_json):214 """215 Segment image with point prompts - MATCHES BACKEND /api/medsam/segment_points216 217 Each point gets its own small segment (converted to small bounding box).218 This matches the backend behavior where points are converted to small boxes.219 220 Args:221 image: PIL Image222 request_json: JSON string with format:223 {224 "points": [[x1, y1], [x2, y2], ...],225 "labels": [1, 0, ...] # 1=foreground, 0=background226 }227 228 Returns:229 JSON string matching backend response format:230 {231 "success": true,232 "masks": [{"mask": [[...]], "confidence": 0.95}, ...],233 "confidences": [0.95, ...],234 "method": "medsam_points_individual"235 }236 """237 try:238 # Parse input239 data = json.loads(request_json)240 points = data.get("points", [])241 labels = data.get("labels", [])242 image_id = data.get("image_id") # Optional: if provided, use precomputed embedding243 244 if not points:245 return json.dumps({'success': False, 'error': 'At least one point is required'})246 247 # Convert PIL to numpy248 image_array = np.array(image)249 H, W = image_array.shape[:2]250 251 # For now, always compute embedding from image (stateless API)252 predictor.set_image(image_array)253 254 # Process each point individually (like backend does)255 box_size = 20 # Small box size for point-based segmentation256 masks_list = []257 confidences_list = []258 259 for i, pt in enumerate(points):260 x, y = pt261 262 # Create a small bounding box centered on the point (matching backend behavior)263 x1 = max(0, x - box_size // 2)264 y1 = max(0, y - box_size // 2)265 x2 = min(W - 1, x + box_size // 2)266 y2 = min(H - 1, y + box_size // 2)267 bbox = np.array([x1, y1, x2, y2])268 269 print(f"Processing point {i+1}/{len(points)}: ({x}, {y}) -> bbox: {bbox.tolist()}")270 271 # Run prediction with box272 masks, scores, logits = predictor.predict(273 point_coords=None,274 point_labels=None,275 box=bbox,276 multimask_output=False277 )278 279 if len(masks) > 0:280 # Take the best mask281 best_idx = np.argmax(scores)282 mask = masks[best_idx]283 score = float(scores[best_idx])284 285 masks_list.append({286 'mask': mask.astype(np.uint8).tolist(),287 'confidence': score288 })289 confidences_list.append(score)290 print(f"Point {i+1} segmentation successful, confidence: {score:.4f}")291 else:292 print(f"Point {i+1} segmentation failed")293 294 if masks_list:295 result = {296 'success': True,297 'masks': masks_list,298 'confidences': confidences_list,299 'method': 'medsam_points_individual'300 }301 else:302 result = {'success': False, 'error': 'All point segmentations failed'}303 304 return json.dumps(result)305 306 except Exception as e:307 import traceback308 return json.dumps({309 'success': False,310 'error': str(e),311 'traceback': traceback.format_exc()312 })313 314 315def segment_box(image, request_json):316 """317 Segment image with a single bounding box - MATCHES BACKEND /api/medsam/segment_box318 319 Args:320 image: PIL Image321 request_json: JSON string with format:322 {323 "bbox": [x1, y1, x2, y2] # Can be array or object with x1,y1,x2,y2324 }325 326 Returns:327 JSON string matching backend response format:328 {329 "success": true,330 "mask": [[...]],331 "confidence": 0.95,332 "method": "medsam_box"333 }334 """335 try:336 # Parse input337 data = json.loads(request_json)338 bbox = data.get("bbox", [])339 340 # Handle both array format [x1,y1,x2,y2] and object format {x1,y1,x2,y2}341 if isinstance(bbox, dict):342 bbox = [bbox.get('x1', 0), bbox.get('y1', 0), bbox.get('x2', 0), bbox.get('y2', 0)]343 344 if not bbox or len(bbox) != 4:345 return json.dumps({'success': False, 'error': 'Valid bounding box required [x1, y1, x2, y2]'})346 347 box = np.array(bbox)348 349 # Convert PIL to numpy350 image_array = np.array(image)351 352 # Stateless: always compute embedding from image353 predictor.set_image(image_array)354 355 # Run prediction with box356 masks, scores, logits = predictor.predict(357 point_coords=None,358 point_labels=None,359 box=box,360 multimask_output=False361 )362 363 if len(masks) > 0:364 best_idx = np.argmax(scores)365 mask = masks[best_idx]366 score = float(scores[best_idx])367 368 result = {369 'success': True,370 'mask': mask.astype(np.uint8).tolist(),371 'confidence': score,372 'method': 'medsam_box'373 }374 else:375 result = {'success': False, 'error': 'Segmentation failed'}376 377 return json.dumps(result)378 379 except Exception as e:380 import traceback381 return json.dumps({382 'success': False,383 'error': str(e),384 'traceback': traceback.format_exc()385 })386 387 388def segment_multiple_boxes(image, request_json):389 """390 Segment image with multiple bounding boxes - MATCHES BACKEND /api/medsam/segment_multiple_boxes391 392 This is the main API endpoint used by the frontend for box-based segmentation.393 394 Args:395 image: PIL Image396 request_json: JSON string with format:397 {398 "bboxes": [399 [x1, y1, x2, y2], # Array format400 {"x1": 10, "y1": 20, "x2": 100, "y2": 200} # Object format (also supported)401 ]402 }403 404 Returns:405 JSON string matching backend response format:406 {407 "success": true,408 "masks": [{"mask": [[...]], "confidence": 0.95}, ...],409 "confidences": [0.95, ...],410 "method": "medsam_multiple_boxes"411 }412 """413 try:414 # Parse input415 data = json.loads(request_json)416 bboxes = data.get("bboxes", [])417 418 if not bboxes:419 return json.dumps({'success': False, 'error': 'At least one bounding box is required'})420 421 # Convert PIL to numpy422 image_array = np.array(image)423 424 # Stateless: always compute embedding from image425 predictor.set_image(image_array)426 427 print(f"Processing {len(bboxes)} boxes for segmentation")428 429 masks_list = []430 confidences_list = []431 432 for i, bbox in enumerate(bboxes):433 # Handle both array format [x1,y1,x2,y2] and object format {x1,y1,x2,y2}434 if isinstance(bbox, dict):435 box = np.array([436 bbox.get('x1', 0),437 bbox.get('y1', 0),438 bbox.get('x2', 0),439 bbox.get('y2', 0)440 ])441 else:442 box = np.array(bbox)443 444 print(f"Processing box {i+1}/{len(bboxes)}: {box.tolist()}")445 446 # Run prediction with box447 masks, scores, logits = predictor.predict(448 point_coords=None,449 point_labels=None,450 box=box,451 multimask_output=False452 )453 454 if len(masks) > 0:455 best_idx = np.argmax(scores)456 mask = masks[best_idx]457 score = float(scores[best_idx])458 459 masks_list.append({460 'mask': mask.astype(np.uint8).tolist(),461 'confidence': score462 })463 confidences_list.append(score)464 print(f"Box {i+1} segmentation successful, confidence: {score:.4f}")465 else:466 print(f"Box {i+1} segmentation failed")467 468 if masks_list:469 result = {470 'success': True,471 'masks': masks_list,472 'confidences': confidences_list,473 'method': 'medsam_multiple_boxes'474 }475 else:476 result = {'success': False, 'error': 'All segmentations failed'}477 478 return json.dumps(result)479 480 except Exception as e:481 import traceback482 return json.dumps({483 'success': False,484 'error': str(e),485 'traceback': traceback.format_exc()486 })487 488 489# =============================================================================490# AUTO MASK GENERATION API (replaces local mask_generator.generate())491# =============================================================================492 493def generate_auto_masks(image, request_json):494 """495 Automatically generate all masks for an image using SAM-H model.496 497 This is equivalent to `mask_generator.generate(img_np)` in enhanced_preprocessing.py498 499 Args:500 image: PIL Image501 request_json: JSON string with optional parameters:502 {503 "points_per_side": 32, # Grid density (default: 32)504 "pred_iou_thresh": 0.88, # IoU threshold (default: 0.88)505 "stability_score_thresh": 0.95, # Stability threshold (default: 0.95)506 "min_mask_region_area": 0 # Minimum mask area (default: 0)507 }508 509 Returns:510 JSON string with format matching SamAutomaticMaskGenerator output:511 {512 "success": true,513 "masks": [514 {515 "segmentation": [[...2D boolean array...]],516 "area": 12345,517 "bbox": [x, y, width, height],518 "predicted_iou": 0.95,519 "point_coords": [[x, y]],520 "stability_score": 0.98,521 "crop_box": [x, y, width, height]522 },523 ...524 ],525 "num_masks": 42,526 "image_size": [height, width]527 }528 """529 try:530 if mask_generator is None:531 return json.dumps({532 'success': False,533 'error': 'MedSAM model not loaded. Please ensure medsam_vit_b.pth is available.',534 'available': False535 })536 537 # Parse optional parameters538 params = {}539 if request_json:540 try:541 params = json.loads(request_json) if request_json.strip() else {}542 except:543 params = {}544 545 # Convert PIL to numpy546 image_array = np.array(image)547 H, W = image_array.shape[:2]548 549 # Optional downscaling to keep masks smaller / faster550 resize_longest = int(params.get("resize_longest", 0) or 0)551 if resize_longest > 0 and max(H, W) > resize_longest:552 scale = resize_longest / float(max(H, W))553 new_w = max(1, int(W * scale))554 new_h = max(1, int(H * scale))555 print(f"Resizing image from {W}x{H} to {new_w}x{new_h} for auto masks...")556 image_array = np.array(Image.fromarray(image_array).resize((new_w, new_h)))557 H, W = image_array.shape[:2]558 559 print(f"Generating automatic masks for image of size {W}x{H}...")560 561 # Generate masks using SAM automatic mask generator562 masks = mask_generator.generate(image_array)563 564 print(f"Generated {len(masks)} masks")565 if len(masks) > 0:566 # Log some stats about the masks567 areas = [m['area'] for m in masks]568 ious = [m['predicted_iou'] for m in masks]569 stabilities = [m['stability_score'] for m in masks]570 print(f" Area range: {min(areas)} - {max(areas)} pixels")571 print(f" IoU range: {min(ious):.3f} - {max(ious):.3f}")572 print(f" Stability range: {min(stabilities):.3f} - {max(stabilities):.3f}")573 else:574 print(" WARNING: No masks generated! This could mean:")575 print(" - Image is too uniform/simple")576 print(" - Thresholds are still too strict")577 print(" - Image size is too small or too large")578 579 # Optionally limit number of masks returned to keep JSON payload reasonable580 max_masks = int(params.get("max_masks", 10))581 if max_masks > 0 and len(masks) > max_masks:582 # Sort by predicted IoU (descending) and keep top-K583 print(f"Limiting masks from {len(masks)} to top {max_masks} by predicted_iou")584 masks = sorted(585 masks,586 key=lambda m: float(m.get("predicted_iou", 0.0)),587 reverse=True,588 )[:max_masks]589 590 print(f"Preparing {len(masks)} masks to return to client...")591 592 # Convert masks to JSON-serializable format593 masks_output = []594 for m in masks:595 mask_data = {596 "segmentation": m["segmentation"].astype(np.uint8).tolist(),597 "area": int(m["area"]),598 "bbox": [int(x) for x in m["bbox"]], # [x, y, width, height]599 "predicted_iou": float(m["predicted_iou"]),600 "point_coords": [601 [int(p[0]), int(p[1])] for p in m["point_coords"]602 ]603 if m["point_coords"] is not None604 else [],605 "stability_score": float(m["stability_score"]),606 "crop_box": [int(x) for x in m["crop_box"]], # [x, y, width, height]607 }608 masks_output.append(mask_data)609 610 result = {611 'success': True,612 'masks': masks_output,613 'num_masks': len(masks_output),614 'image_size': [H, W]615 }616 617 print(f"Auto mask generation complete: {len(masks_output)} masks")618 return json.dumps(result)619 620 except Exception as e:621 import traceback622 return json.dumps({623 'success': False,624 'error': str(e),625 'traceback': traceback.format_exc()626 })627 628 629def check_auto_mask_status():630 """631 Check if automatic mask generation is available632 """633 return json.dumps({634 'available': mask_generator is not None,635 'model': MODEL_FILENAME if mask_generator else None,636 'model_type': MODEL_TYPE,637 'device': str(device)638 })639 640 641# =============================================================================642# LEGACY API FUNCTIONS (kept for backwards compatibility with test scripts)643# =============================================================================644 645def segment_with_points_legacy(image, points_json):646 """647 Legacy API - Segment with point prompts using true point-based segmentation648 649 Args:650 points_json: JSON string with format:651 {652 "coords": [[x1, y1], [x2, y2], ...],653 "labels": [1, 0, ...],654 "multimask_output": true/false655 }656 """657 try:658 points_data = json.loads(points_json)659 coords = np.array(points_data["coords"])660 labels = np.array(points_data["labels"])661 multimask_output = points_data.get("multimask_output", True)662 663 image_array = np.array(image)664 predictor.set_image(image_array)665 666 masks, scores, logits = predictor.predict(667 point_coords=coords,668 point_labels=labels,669 multimask_output=multimask_output670 )671 672 masks_list = []673 scores_list = []674 675 for i, (mask, score) in enumerate(zip(masks, scores)):676 mask_uint8 = (mask * 255).astype(np.uint8)677 mask_image = Image.fromarray(mask_uint8)678 buffer = io.BytesIO()679 mask_image.save(buffer, format='PNG')680 mask_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')681 682 masks_list.append({683 'mask_base64': mask_base64,684 'mask_shape': mask.shape,685 'mask_data': mask.tolist()686 })687 scores_list.append(float(score))688 689 return json.dumps({690 'success': True,691 'masks': masks_list,692 'scores': scores_list,693 'num_masks': len(masks_list)694 })695 696 except Exception as e:697 return json.dumps({'success': False, 'error': str(e)})698 699 700def segment_with_box_legacy(image, box_json):701 """702 Legacy API - Segment with box prompt703 704 Args:705 box_json: JSON string with format:706 {"box": [x1, y1, x2, y2], "multimask_output": false}707 """708 try:709 box_data = json.loads(box_json)710 box = np.array(box_data["box"])711 multimask_output = box_data.get("multimask_output", False)712 713 image_array = np.array(image)714 predictor.set_image(image_array)715 716 masks, scores, logits = predictor.predict(717 point_coords=None,718 point_labels=None,719 box=box,720 multimask_output=multimask_output721 )722 723 masks_list = []724 scores_list = []725 726 for i, (mask, score) in enumerate(zip(masks, scores)):727 mask_uint8 = (mask * 255).astype(np.uint8)728 mask_image = Image.fromarray(mask_uint8)729 buffer = io.BytesIO()730 mask_image.save(buffer, format='PNG')731 mask_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')732 733 masks_list.append({734 'mask_base64': mask_base64,735 'mask_shape': mask.shape,736 'mask_data': mask.tolist()737 })738 scores_list.append(float(score))739 740 return json.dumps({741 'success': True,742 'masks': masks_list,743 'scores': scores_list,744 'num_masks': len(masks_list),745 'box': box.tolist()746 })747 748 except Exception as e:749 import traceback750 return json.dumps({751 'success': False,752 'error': str(e),753 'traceback': traceback.format_exc()754 })755 756 757def segment_simple(image, x, y, label=1, multimask=True):758 """Simple single-point segmentation for Gradio UI"""759 try:760 points_json = json.dumps({761 "coords": [[int(x), int(y)]],762 "labels": [int(label)],763 "multimask_output": multimask764 })765 766 result_json = segment_with_points_legacy(image, points_json)767 result = json.loads(result_json)768 769 if not result['success']:770 return None, f"Error: {result['error']}"771 772 best_idx = np.argmax(result['scores'])773 best_mask_base64 = result['masks'][best_idx]['mask_base64']774 best_score = result['scores'][best_idx]775 776 mask_bytes = base64.b64decode(best_mask_base64)777 mask_image = Image.open(io.BytesIO(mask_bytes))778 779 return mask_image, f"Score: {best_score:.4f}"780 781 except Exception as e:782 return None, f"Error: {str(e)}"783 784 785# =============================================================================786# GRADIO INTERFACE787# =============================================================================788 789with gr.Blocks(title="MedSAM Inference API") as demo:790 gr.Markdown("# 🏥 MedSAM Inference API")791 gr.Markdown("Point and box-based segmentation using Fine-Tuned MedSAM")792 gr.Markdown("**API-compatible with Dense-Captioning-Toolkit backend**")793 794 with gr.Tabs():795 # Tab 1: Backend-Compatible API (Points)796 with gr.Tab("Segment Points (Backend API)"):797 gr.Markdown("""798 ## Point-based Segmentation - Backend Compatible799 800 **Matches `/api/medsam/segment_points`**801 802 Each point is converted to a small bounding box for segmentation.803 804 **Input Format:**805 ```json806 {807 "points": [[x1, y1], [x2, y2], ...],808 "labels": [1, 0, ...]809 }810 ```811 812 **Output Format (matches backend):**813 ```json814 {815 "success": true,816 "masks": [{"mask": [[...]], "confidence": 0.95}, ...],817 "confidences": [0.95, ...],818 "method": "medsam_points_individual"819 }820 ```821 """)822 823 with gr.Row():824 with gr.Column():825 points_image = gr.Image(type="pil", label="Input Image")826 points_json_input = gr.Textbox(827 label="Request JSON",828 placeholder='{"points": [[100, 150], [200, 200]], "labels": [1, 1]}',829 lines=3830 )831 points_button = gr.Button("Segment Points", variant="primary")832 833 with gr.Column():834 points_output = gr.Textbox(label="Result JSON", lines=15)835 836 points_button.click(837 fn=segment_points,838 inputs=[points_image, points_json_input],839 outputs=points_output,840 api_name="segment_points"841 )842 843 # Tab 2: Backend-Compatible API (Multiple Boxes)844 with gr.Tab("Segment Multiple Boxes (Backend API)"):845 gr.Markdown("""846 ## Multiple Box Segmentation - Backend Compatible847 848 **Matches `/api/medsam/segment_multiple_boxes`** (main frontend API)849 850 **Input Format:**851 ```json852 {853 "bboxes": [854 [x1, y1, x2, y2],855 {"x1": 10, "y1": 20, "x2": 100, "y2": 200}856 ]857 }858 ```859 860 **Output Format (matches backend):**861 ```json862 {863 "success": true,864 "masks": [{"mask": [[...]], "confidence": 0.95}, ...],865 "confidences": [0.95, ...],866 "method": "medsam_multiple_boxes"867 }868 ```869 """)870 871 with gr.Row():872 with gr.Column():873 multi_box_image = gr.Image(type="pil", label="Input Image")874 multi_box_json = gr.Textbox(875 label="Request JSON",876 placeholder='{"bboxes": [[100, 100, 300, 300], [400, 400, 600, 600]]}',877 lines=3878 )879 multi_box_button = gr.Button("Segment Multiple Boxes", variant="primary")880 881 with gr.Column():882 multi_box_output = gr.Textbox(label="Result JSON", lines=15)883 884 multi_box_button.click(885 fn=segment_multiple_boxes,886 inputs=[multi_box_image, multi_box_json],887 outputs=multi_box_output,888 api_name="segment_multiple_boxes"889 )890 891 # Tab 3: Backend-Compatible API (Single Box)892 with gr.Tab("Segment Box (Backend API)"):893 gr.Markdown("""894 ## Single Box Segmentation - Backend Compatible895 896 **Matches `/api/medsam/segment_box`**897 898 **Input Format:**899 ```json900 {901 "bbox": [x1, y1, x2, y2]902 }903 ```904 905 **Output Format (matches backend):**906 ```json907 {908 "success": true,909 "mask": [[...]],910 "confidence": 0.95,911 "method": "medsam_box"912 }913 ```914 """)915 916 with gr.Row():917 with gr.Column():918 box_image = gr.Image(type="pil", label="Input Image")919 box_json_input = gr.Textbox(920 label="Request JSON",921 placeholder='{"bbox": [100, 100, 300, 300]}',922 lines=3923 )924 box_button = gr.Button("Segment Box", variant="primary")925 926 with gr.Column():927 box_output = gr.Textbox(label="Result JSON", lines=15)928 929 box_button.click(930 fn=segment_box,931 inputs=[box_image, box_json_input],932 outputs=box_output,933 api_name="segment_box"934 )935 936 # Tab 4: Legacy API (for test scripts)937 with gr.Tab("Legacy API"):938 gr.Markdown("""939 ## Legacy API (for backwards compatibility)940 941 Original API format with `coords`, `mask_data`, `scores`, etc.942 Use if you have existing scripts using the old format.943 """)944 945 with gr.Row():946 with gr.Column():947 legacy_image = gr.Image(type="pil", label="Input Image")948 legacy_points = gr.Textbox(949 label="Points JSON (Legacy Format)",950 placeholder='{"coords": [[100, 150]], "labels": [1], "multimask_output": true}',951 lines=3952 )953 legacy_button = gr.Button("Run Segmentation (Legacy)", variant="secondary")954 955 with gr.Column():956 legacy_output = gr.Textbox(label="Result JSON", lines=15)957 958 legacy_button.click(959 fn=segment_with_points_legacy,960 inputs=[legacy_image, legacy_points],961 outputs=legacy_output,962 api_name="segment_with_points" # Keep old API name for compatibility963 )964 965 gr.Markdown("---")966 967 with gr.Row():968 with gr.Column():969 legacy_box_image = gr.Image(type="pil", label="Input Image")970 legacy_box_json = gr.Textbox(971 label="Box JSON (Legacy Format)",972 placeholder='{"box": [100, 100, 300, 300], "multimask_output": false}',973 lines=3974 )975 legacy_box_button = gr.Button("Run Box Segmentation (Legacy)", variant="secondary")976 977 with gr.Column():978 legacy_box_output = gr.Textbox(label="Result JSON", lines=15)979 980 legacy_box_button.click(981 fn=segment_with_box_legacy,982 inputs=[legacy_box_image, legacy_box_json],983 outputs=legacy_box_output,984 api_name="segment_with_box" # Keep old API name for compatibility985 )986 987 # Tab 5: Auto Mask Generation (for preprocessing)988 with gr.Tab("Auto Mask Generation"):989 gr.Markdown("""990 ## Automatic Mask Generation (MedSAM)991 992 **Replaces `mask_generator.generate(img_np)` in preprocessing pipeline**993 994 Uses MedSAM (ViT-B) model with `SamAutomaticMaskGenerator` to automatically 995 segment all objects in an image. This is used for initial preprocessing 996 of scientific/medical images.997 998 Uses the same `medsam_vit_b.pth` model as interactive segmentation.999 1000 **Output Format:**1001 ```json1002 {1003 "success": true,1004 "masks": [1005 {1006 "segmentation": [[...2D array...]],1007 "area": 12345,1008 "bbox": [x, y, width, height],1009 "predicted_iou": 0.95,1010 "point_coords": [[x, y]],1011 "stability_score": 0.98,1012 "crop_box": [x, y, width, height]1013 }1014 ],1015 "num_masks": 421016 }1017 ```1018 """)1019 1020 with gr.Row():1021 with gr.Column():1022 auto_image = gr.Image(type="pil", label="Input Image")1023 auto_params = gr.Textbox(1024 label="Parameters (optional)",1025 placeholder='{"points_per_side": 32, "pred_iou_thresh": 0.88}',1026 lines=21027 )1028 with gr.Row():1029 auto_button = gr.Button("Generate All Masks", variant="primary")1030 status_button = gr.Button("Check Status", variant="secondary")1031 1032 with gr.Column():1033 auto_output = gr.Textbox(label="Result JSON", lines=20)1034 status_output = gr.Textbox(label="Status", lines=3)1035 1036 auto_button.click(1037 fn=generate_auto_masks,1038 inputs=[auto_image, auto_params],1039 outputs=auto_output,1040 api_name="generate_auto_masks"1041 )1042 1043 status_button.click(1044 fn=check_auto_mask_status,1045 inputs=[],1046 outputs=status_output,1047 api_name="check_auto_mask_status"1048 )1049 1050 # Tab 6: Encode Image (for embedding storage)1051 with gr.Tab("Encode Image"):1052 gr.Markdown("""1053 ## Image Encoding API1054 1055 **Encodes image using SAM image encoder and saves embedding to Supabase**1056 1057 This endpoint is used during preprocessing to compute and store image embeddings1058 once per image. Later segmentation calls can use these precomputed embeddings1059 for faster inference (no need to recompute embeddings on each API call).1060 1061 **Input Format:**1062 ```json1063 {1064 "image_id": "uuid-string" # Required: image ID from database1065 }1066 ```1067 1068 **Output Format:**1069 ```json1070 {1071 "success": true,1072 "message": "Embedding saved successfully for image_id=...",1073 "image_id": "uuid-string",1074 "embedding_shape": [1, 256, 64, 64]1075 }1076 ```1077 1078 **Note:** Requires Supabase credentials (SUPABASE_URL and SUPABASE_KEY environment variables)1079 """)1080 1081 with gr.Row():1082 with gr.Column():1083 encode_image_input = gr.Image(type="pil", label="Input Image")1084 encode_json_input = gr.Textbox(1085 label="Request JSON",1086 placeholder='{"image_id": "123e4567-e89b-12d3-a456-426614174000"}',1087 lines=21088 )1089 encode_button = gr.Button("Encode Image", variant="primary")1090 1091 with gr.Column():1092 encode_output = gr.Textbox(label="Result JSON", lines=10)1093 1094 encode_button.click(1095 fn=encode_image,1096 inputs=[encode_image_input, encode_json_input],1097 outputs=encode_output,1098 api_name="encode_image"1099 )1100 1101 # Tab 7: Simple UI Interface1102 with gr.Tab("Simple Interface"):1103 gr.Markdown("## Click-based Segmentation")1104 gr.Markdown("Enter X, Y coordinates to segment")1105 1106 with gr.Row():1107 with gr.Column():1108 simple_image = gr.Image(type="pil", label="Input Image")1109 with gr.Row():1110 simple_x = gr.Number(label="X Coordinate", value=100)1111 simple_y = gr.Number(label="Y Coordinate", value=100)1112 with gr.Row():1113 simple_label = gr.Radio(1114 choices=[1, 0],1115 value=1,1116 label="Point Label (1=foreground, 0=background)"1117 )1118 simple_multimask = gr.Checkbox(1119 label="Multiple Masks",1120 value=True1121 )1122 simple_button = gr.Button("Segment", variant="primary")1123 1124 with gr.Column():1125 simple_mask = gr.Image(label="Output Mask")1126 simple_info = gr.Textbox(label="Info")1127 1128 simple_button.click(1129 fn=segment_simple,1130 inputs=[simple_image, simple_x, simple_y, simple_label, simple_multimask],1131 outputs=[simple_mask, simple_info]1132 )1133 1134 gr.Markdown("""1135 ---1136 ### 📡 API Usage from Python (Backend-Compatible)1137 1138 ```python1139 from gradio_client import Client, handle_file1140 import json1141 1142 client = Client("Aniketg6/medsam-inference")1143 1144 # Point-based segmentation (matches backend format)1145 result = client.predict(1146 image=handle_file("image.jpg"),1147 request_json=json.dumps({1148 "points": [[150, 200], [300, 400]],1149 "labels": [1, 1]1150 }),1151 api_name="/segment_points"1152 )1153 1154 # Multiple box segmentation (main frontend API)1155 result = client.predict(1156 image=handle_file("image.jpg"),1157 request_json=json.dumps({1158 "bboxes": [[100, 100, 300, 300], [400, 400, 600, 600]]1159 }),1160 api_name="/segment_multiple_boxes"1161 )1162 1163 # Parse response1164 data = json.loads(result)1165 print(f"Success: {data['success']}")1166 print(f"Masks: {len(data['masks'])}")1167 print(f"Confidences: {data['confidences']}")1168 print(f"Method: {data['method']}")1169 ```1170 """)1171 1172# Launch1173if __name__ == "__main__":1174 demo.launch(1175 server_name="0.0.0.0",1176 server_port=7860,1177 share=False,1178 show_error=True1179 )1180