Upload folder using huggingface_hub

app.py

@@ -60,9 +60,13 @@ CHECKPOINTS = {
 HF_REPO_ID = os.getenv("HF_REPO_ID", "tp53/oncoseg")
 
 
+# Flag to track if we're using fallback mode
+USE_FALLBACK = False
+
+
 def get_model(checkpoint: str = "brain"):
-    """Load or retrieve cached model."""
-    global MODELS
+    """Load or retrieve cached model. Falls back to simple segmentation if SAM3 unavailable."""
+    global MODELS, USE_FALLBACK
 
     if checkpoint not in MODELS:
         logger.info(f"Loading model: {checkpoint}")
@@ -89,11 +93,45 @@ def get_model(checkpoint: str = "brain"):
             MODELS[checkpoint] = model
             logger.info(f"Model {checkpoint} loaded on {DEVICE}")
 
+        except ImportError as e:
+            logger.warning(f"SAM3 not available, using fallback segmentation: {e}")
+            USE_FALLBACK = True
+            MODELS[checkpoint] = None
         except Exception as e:
             logger.error(f"Failed to load model {checkpoint}: {e}")
-            raise
+            USE_FALLBACK = True
+            MODELS[checkpoint] = None
+
+    return MODELS.get(checkpoint)
+
+
+def fallback_segment(slice_2d: np.ndarray):
+    """
+    Simple intensity-based segmentation fallback when SAM3 is not available.
+    Works well for FLAIR MRI where tumors appear hyperintense.
+    """
+    from skimage.filters import threshold_otsu
+    from skimage.morphology import binary_opening, binary_closing, disk
+
+    # Normalize
+    vmin, vmax = slice_2d.min(), slice_2d.max()
+    if vmax - vmin < 1e-8:
+        return np.zeros_like(slice_2d, dtype=np.uint8)
+
+    normalized = (slice_2d - vmin) / (vmax - vmin)
+
+    # Use percentile threshold (top 15% intensity = potential tumor)
+    threshold = np.percentile(normalized, 85)
+    mask = (normalized > threshold).astype(np.uint8)
+
+    # Morphological cleanup
+    try:
+        mask = binary_opening(mask, disk(2))
+        mask = binary_closing(mask, disk(3))
+    except:
+        pass
 
-    return MODELS[checkpoint]
+    return mask.astype(np.uint8)
 
 
 def preprocess_slice(slice_2d: np.ndarray, target_size: int = 1024) -> torch.Tensor:
@@ -204,44 +242,53 @@ def _segment_slice_impl(
                 "error": f"Slice index {slice_idx} out of range [0, {volume.shape[0]})",
             }
 
-        # Load model
-        model = get_model(checkpoint)
-
-        # Extract and preprocess slice
+        # Extract slice
         slice_2d = volume[slice_idx]
         original_shape = slice_2d.shape
-        slice_tensor = preprocess_slice(slice_2d)  # (1, 3, 1024, 1024) tensor on DEVICE
 
-        # Create full-image bounding box prompt (auto-segment entire image)
-        # Format: [x_min, y_min, x_max, y_max] in pixel coordinates
-        target_size = slice_tensor.shape[-1]  # 1024
-        input_boxes = torch.tensor(
-            [[0, 0, target_size, target_size]], dtype=torch.float32, device=DEVICE
-        )
+        # Load model (may return None if fallback mode)
+        model = get_model(checkpoint)
 
-        # Run inference with text prompt for grounding
-        with torch.no_grad():
-            outputs = model(
-                pixel_values=slice_tensor,
-                input_boxes=input_boxes,
-                text_prompt=text_prompt,
+        if model is None or USE_FALLBACK:
+            # Use fallback segmentation
+            logger.info("Using fallback segmentation (SAM3 not available)")
+            mask = fallback_segment(slice_2d)
+            backend = "fallback"
+        else:
+            # Use SAM3 model
+            slice_tensor = preprocess_slice(
+                slice_2d
+            )  # (1, 3, 1024, 1024) tensor on DEVICE
+
+            # Create full-image bounding box prompt (auto-segment entire image)
+            # Format: [x_min, y_min, x_max, y_max] in pixel coordinates
+            target_size = slice_tensor.shape[-1]  # 1024
+            input_boxes = torch.tensor(
+                [[0, 0, target_size, target_size]], dtype=torch.float32, device=DEVICE
             )
 
-        # Extract mask from SAM3 output
-        # SAM3 returns a dict with 'pred_masks' key, shape (B, 1, H, W)
-        if isinstance(outputs, dict) and "pred_masks" in outputs:
-            pred_mask = outputs["pred_masks"][0, 0].cpu().numpy()
-        elif hasattr(outputs, "pred_masks"):
-            pred_mask = outputs.pred_masks[0, 0].cpu().numpy()
-        else:
-            # Fallback: try to extract from tuple/list
-            logger.warning(f"Unexpected output type: {type(outputs)}")
-            pred_mask = np.zeros((target_size, target_size))
+            # Run inference with text prompt for grounding
+            with torch.no_grad():
+                outputs = model(
+                    pixel_values=slice_tensor,
+                    input_boxes=input_boxes,
+                    text_prompt=text_prompt,
+                )
 
-        # Resize mask back to original shape
-        import cv2
+            # Extract mask from SAM3 output
+            # SAM3 returns a dict with 'pred_masks' key, shape (B, 1, H, W)
+            if isinstance(outputs, dict) and "pred_masks" in outputs:
+                pred_mask = outputs["pred_masks"][0, 0].cpu().numpy()
+            elif hasattr(outputs, "pred_masks"):
+                pred_mask = outputs.pred_masks[0, 0].cpu().numpy()
+            else:
+                # Fallback: try to extract from tuple/list
+                logger.warning(f"Unexpected output type: {type(outputs)}")
+                pred_mask = np.zeros((target_size, target_size))
 
-        mask = cv2.resize(pred_mask, (original_shape[1], original_shape[0]))
+            # Resize mask back to original shape
+            mask = cv2.resize(pred_mask, (original_shape[1], original_shape[0]))
+            backend = "sam3"
 
         # Threshold to binary
         mask = (mask > 0.5).astype(np.uint8)
@@ -260,6 +307,7 @@ def _segment_slice_impl(
 
         return {
             "success": True,
+            "backend": backend,
             "mask_b64": mask_b64,
             "mask_shape": list(mask.shape),
             "contours": contours,
@@ -310,8 +358,9 @@ def _segment_volume_impl(
 
         logger.info(f"Loaded volume shape: {volume.shape}")
 
-        # Load model
+        # Load model (may return None if fallback mode)
         model = get_model(checkpoint)
+        use_fallback = model is None or USE_FALLBACK
 
         num_slices = volume.shape[0]
         all_contours = {}
@@ -326,31 +375,37 @@ def _segment_volume_impl(
             if skip_empty and slice_2d.max() - slice_2d.min() < 0.01:
                 continue
 
-            slice_tensor = preprocess_slice(slice_2d, target_size)
-
-            # Create full-image bounding box
-            input_boxes = torch.tensor(
-                [[0, 0, target_size, target_size]], dtype=torch.float32, device=DEVICE
-            )
+            if use_fallback:
+                # Use fallback segmentation
+                mask = fallback_segment(slice_2d)
+            else:
+                slice_tensor = preprocess_slice(slice_2d, target_size)
 
-            with torch.no_grad():
-                outputs = model(
-                    pixel_values=slice_tensor,
-                    input_boxes=input_boxes,
-                    text_prompt=text_prompt,
+                # Create full-image bounding box
+                input_boxes = torch.tensor(
+                    [[0, 0, target_size, target_size]],
+                    dtype=torch.float32,
+                    device=DEVICE,
                 )
 
-            # Extract mask from SAM3 output
-            if isinstance(outputs, dict) and "pred_masks" in outputs:
-                pred_mask = outputs["pred_masks"][0, 0].cpu().numpy()
-            elif hasattr(outputs, "pred_masks"):
-                pred_mask = outputs.pred_masks[0, 0].cpu().numpy()
-            else:
-                continue  # Skip if no valid output
-
-            # Resize to original shape and threshold
-            mask = cv2.resize(pred_mask, (original_shape[1], original_shape[0]))
-            mask = (mask > 0.5).astype(np.uint8)
+                with torch.no_grad():
+                    outputs = model(
+                        pixel_values=slice_tensor,
+                        input_boxes=input_boxes,
+                        text_prompt=text_prompt,
+                    )
+
+                # Extract mask from SAM3 output
+                if isinstance(outputs, dict) and "pred_masks" in outputs:
+                    pred_mask = outputs["pred_masks"][0, 0].cpu().numpy()
+                elif hasattr(outputs, "pred_masks"):
+                    pred_mask = outputs.pred_masks[0, 0].cpu().numpy()
+                else:
+                    continue  # Skip if no valid output
+
+                # Resize to original shape and threshold
+                mask = cv2.resize(pred_mask, (original_shape[1], original_shape[0]))
+                mask = (mask > 0.5).astype(np.uint8)
 
             if mask.sum() >= min_area:
                 mask = keep_largest_component(mask)

requirements.txt

@@ -17,10 +17,7 @@ opencv-python-headless>=4.8.0
 huggingface_hub>=0.19.0
 spaces>=0.19.0
 
-# SAM3 - Install from Facebook Research
-git+https://github.com/facebookresearch/sam3.git
-
-# Additional SAM3 Dependencies
+# SAM3 Dependencies (SAM3 itself installed separately if available)
 einops>=0.7.0
 timm>=0.9.0
 peft>=0.7.0