project-lighter
/

ct_fm_feature_extractor

@@ -18,15 +18,13 @@ The backbone is based on a SegResNet, a 3D U-Net variant. If you want to just lo
 ## Running instructions
-# Whole Body Segmentation Inference
 This notebook demonstrates how to:
-1. Load a pre-trained whole body segmentation model from HuggingFace Hub
 2. Set up preprocessing and postprocessing pipelines
-3. Perform sliding window inference on CT volumes
-4. Save the segmentation results
-The model segments 118 different anatomical structures from CT scans.
 ## Setup
 Install requirements and import necessary packages
@@ -44,55 +42,27 @@ Install requirements and import necessary packages
 ```python
 # Imports
 import torch
-from lighter_zoo import SegResNet
 from monai.transforms import (
     Compose, LoadImage, EnsureType, Orientation,
-    ScaleIntensityRange, CropForeground, Invert,
-    Activations, AsDiscrete, KeepLargestConnectedComponent,
-    SaveImage
 )
 from monai.inferers import SlidingWindowInferer
 ```
-    Note: you may need to restart the kernel to use updated packages.
 ## Load Model
 Download and initialize the pre-trained model from HuggingFace Hub
 ```python
 # Load pre-trained model
-model = SegResNet.from_pretrained(
-    "project-lighter/whole_body_segmentation",
-    force_download=True
-)
-```
-    config.json:   0%|          | 0.00/162 [00:00<?, ?B/s]
-    model.safetensors:   0%|          | 0.00/349M [00:00<?, ?B/s]
-## Configure Inference
-Set up sliding window inference for processing large volumes
-```python
-# Configure sliding window inference
-inferer = SlidingWindowInferer(
-    roi_size=[96, 160, 160],  # Size of patches to process
-    sw_batch_size=2,          # Number of windows to process in parallel
-    overlap=0.625,            # Overlap between windows (reduces boundary artifacts)
-    mode="gaussian"           # Gaussian weighting for overlap regions
 )
 ```
 ## Setup Processing Pipelines
-Define preprocessing and postprocessing transforms
 ```python
@@ -111,24 +81,13 @@ preprocess = Compose([
     ),
     CropForeground()    # Remove background to reduce computation
 ])
-# Postprocessing pipeline
-postprocess = Compose([
-    Activations(softmax=True),              # Apply softmax to get probabilities
-    AsDiscrete(argmax=True, dtype=torch.int32),  # Convert to class labels
-    KeepLargestConnectedComponent(),        # Remove small disconnected regions
-    Invert(transform=preprocess),           # Restore original space
-    # Save the result
-    SaveImage(output_dir="./segmentations")
-])
 ```
-    /home/suraj/miniconda3/lib/python3.10/site-packages/monai/utils/deprecate_utils.py:321: FutureWarning: monai.transforms.croppad.array CropForeground.__init__:allow_smaller: Current default value of argument `allow_smaller=True` has been deprecated since version 1.2. It will be changed to `allow_smaller=False` in version 1.5.
-      warn_deprecated(argname, msg, warning_category)
 ## Run Inference
-Process an input CT scan and generate segmentation
 ```python
@@ -140,19 +99,36 @@ input_tensor = preprocess(input_path)
 # Run inference
 with torch.no_grad():
-    output = inferer(input_tensor.unsqueeze(dim=0), model)[0]
-# Copy metadata from input
-output.applied_operations = input_tensor.applied_operations
-output.affine = input_tensor.affine
-# Postprocess and save result
-result = postprocess(output[0])
-print("✅ Segmentation completed and saved")
 ```
-    2025-01-16 18:41:57,674 INFO image_writer.py:197 - writing: /home/suraj/Repositories/lighter-ct-fm/semantic-search-app/assets/segmentations/0/0_trans.nii.gz
-    ✅ Segmentation completed and saved

 ## Running instructions
+# CT-FM Feature Extractor
 This notebook demonstrates how to:
+1. Load a SSL pre-trained model
 2. Set up preprocessing and postprocessing pipelines
+3. Perform inference on CT volumes
+4. Plot distribution of features extracted
 ## Setup
 Install requirements and import necessary packages
 ```python
 # Imports
 import torch
+from lighter_zoo import SegResEncoder
 from monai.transforms import (
     Compose, LoadImage, EnsureType, Orientation,
+    ScaleIntensityRange, CropForeground
 )
 from monai.inferers import SlidingWindowInferer
 ```
 ## Load Model
 Download and initialize the pre-trained model from HuggingFace Hub
 ```python
 # Load pre-trained model
+model = SegResEncoder.from_pretrained(
+    "project-lighter/ct_fm_feature_extractor"
 )
 ```
 ## Setup Processing Pipelines
+Define preprocessing transforms
 ```python
     ),
     CropForeground()    # Remove background to reduce computation
 ])
 ```
+    monai.transforms.croppad.array CropForeground.__init__:allow_smaller: Current default value of argument `allow_smaller=True` has been deprecated since version 1.2. It will be changed to `allow_smaller=False` in version 1.5.
 ## Run Inference
+Process an input CT scan and extract features
 ```python
 # Run inference
 with torch.no_grad():
+    output = model(input_tensor.unsqueeze(0))[-1]
+    # Average pooling compressed the feature vector across all patches. If this is not desired, remove this line and
+    # use the output tensor directly which will give you the feature maps in a low-dimensional space.
+    avg_output = torch.nn.functional.adaptive_avg_pool3d(output, 1).squeeze()
+print("✅ Feature extraction completed")
+print(f"Output shape: {avg_output.shape}")
 ```
+    ✅ Feature extraction completed
+    Output shape: torch.Size([512])
+```python
+# Plot distribution of features
+import matplotlib.pyplot as plt
+_ = plt.hist(avg_output.cpu().numpy(), bins=100)
+```
+![png](ct_fm_feature_extractor_files/ct_fm_feature_extractor_10_0.png)
+```python
+```