Update with README.md, requirements.txt, and inference.py for Inference Endpoint

README.md

@@ -0,0 +1,121 @@
+
+---
+pretty_name: SAM Brain Tumor Segmentation
+tagged_ids:
+- sam-brain-tumor-segmentation
+pipeline_tag: image-segmentation
+library_name: transformers
+license: mit
+---
+
+# SAM Brain Tumor Segmentation Model
+
+This model is a fine-tuned [Segment Anything Model (SAM)](https://huggingface.co/facebook/sam-vit-base) for brain tumor segmentation from medical imaging data. It was trained using a simulated dataset of 2D slices derived from 3D NIfTI (.nii.gz) images and their corresponding segmentation masks.
+
+## Model Description
+
+The original SAM model is a powerful general-purpose image segmentation model. This fine-tuned version specializes in identifying brain tumors, leveraging the prompt-based segmentation capabilities of SAM. The model is prompted with bounding boxes around the tumor regions (derived from ground truth masks during training) to generate precise segmentation masks.
+
+### Training Details
+
+- **Base Model**: `facebook/sam-vit-base`
+- **Dataset**: Simulated 2D axial slices from 3D NIfTI images, normalized to 0-1 range.
+- **Image Preprocessing**: Grayscale images were duplicated across 3 channels to match SAM's expected input. Bounding box prompts were generated from ground truth masks.
+- **Loss Functions**: Binary Cross-Entropy (BCE) Loss and Dice Loss.
+- **Optimizer**: AdamW with a learning rate of 1e-5.
+- **Epochs**: 5
+- **Average Dice Score on Validation Set**: 0.9756 (on simulated data)
+
+## Usage
+
+To use this model for inference, you can load it with the `transformers` library and provide an image along with a bounding box prompt for the region of interest. The model will then predict a segmentation mask.
+
+```python
+from transformers import SamModel, SamProcessor
+from PIL import Image
+import torch
+import numpy as np
+
+# Load the fine-tuned model and processor
+processor = SamProcessor.from_pretrained("Lorenzob/sam-brain-tumor-segmentation")
+model = SamModel.from_pretrained("Lorenzob/sam-brain-tumor-segmentation")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+# Example: Create a dummy image (replace with your actual medical image)
+# This should be a 2D grayscale image, then converted to 3 channels.
+# For a real image, load it and ensure it's normalized 0-1 and uint8 or float.
+image_size = 256 # Example size
+dummy_image_data = np.random.rand(image_size, image_size) * 255
+dummy_image = Image.fromarray(dummy_image_data.astype(np.uint8)).convert("RGB")
+
+# Example: Define a bounding box for the tumor region (x_min, y_min, x_max, y_max)
+# In a real scenario, this bounding box would be provided by an expert or a detection model.
+input_boxes = [[100, 100, 200, 200]] # Example bounding box coordinates
+
+# Preprocess the image and bounding box
+inputs = processor(dummy_image, input_boxes=input_boxes, return_tensors="pt").to(device)
+
+# Perform inference
+with torch.no_grad():
+    outputs = model(**inputs, multimask_output=False)
+
+# Post-process the predicted mask
+masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
+
+# The output `masks` is a list of dictionaries. Each dict contains 'segmentation'.
+# For simplicity, let's take the first mask (assuming multimask_output=False)
+predicted_mask = masks[0]['segmentation'].squeeze().numpy() # Shape (H, W)
+
+print("Predicted mask shape:", predicted_mask.shape)
+# You can visualize 'predicted_mask' using matplotlib or other image libraries.
+# For example:
+# import matplotlib.pyplot as plt
+# plt.imshow(predicted_mask, cmap='gray')
+# plt.title('Predicted Segmentation Mask')
+# plt.show()
+```
+
+## Inference Endpoint Configuration (Optional)
+
+If you wish to deploy this model as an Inference Endpoint on Hugging Face, here's a sample configuration you might use in your `README.md` (or directly in the UI):
+
+```yaml
+widget:
+- src: "app.py"
+  example_title: "Brain Tumor Segmentation Example"
+  inputs:
+  - filename: "image.png"
+    image: https://huggingface.co/datasets/huggingface/sample-images/resolve/main/segmentation_image_input.png
+    input_boxes: [[100, 100, 200, 200]]
+
+--- # Optional section for specific endpoint settings
+
+parameters:
+  do_normalize: false # Assuming inputs are already normalized 0-1
+  do_rescale: false   # Assuming inputs are already scaled correctly
+  multimask_output: false # For single best mask output
+
+# Example of specific hardware/software config for advanced users
+# inference:
+#   accelerator: cuda
+#   container: pytorch_latest
+#   hardware: gpu_small
+#   task: image-segmentation
+```
+
+**Note**: The example image and `input_boxes` in the YAML configuration are placeholders. For a real medical image endpoint, you would provide a relevant example image and a bounding box corresponding to a tumor within that image.
+
+## Limitations
+
+- The model was fine-tuned on a simulated dataset. Its performance on real, diverse clinical data may vary and needs further rigorous validation.
+- The model relies on a bounding box prompt. Its accuracy is highly dependent on the quality and precision of the provided bounding box.
+- Currently, the model handles 2D slices. Adaptation for full 3D volume segmentation would require further development.
+
+## Future Work
+
+- Evaluate and fine-tune the model on large, real-world medical imaging datasets (e.g., BraTS, TCIA).
+- Explore methods for automatic bounding box generation for tumor regions.
+- Extend the model to handle 3D medical images directly.
+- Implement quantitative metrics (e.g., IoU, Hausdorff Distance) during evaluation with real data.

inference.py

@@ -0,0 +1,82 @@
+
+import torch
+from transformers import SamModel, SamProcessor
+from PIL import Image
+import numpy as np
+import io
+import base64
+import json
+
+class InferenceHandler:
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = SamModel.from_pretrained("./sam_brain_tumor_model").to(self.device)
+        self.processor = SamProcessor.from_pretrained("./sam_brain_tumor_model")
+
+    def preprocess(self, request_body):
+        # Expect request_body to be a JSON string with 'image' (base64) and 'boxes' (list of list of floats)
+        data = json.loads(request_body)
+
+        # Decode image from base64
+        image_bytes = base64.b64decode(data['image'])
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+        # Get bounding boxes
+        input_boxes = data.get('boxes', [])
+        # Ensure boxes are in the expected format (list of list of 4 floats)
+        input_boxes = [[float(coord) for coord in box] for box in input_boxes]
+
+        # Prepare inputs for the model
+        inputs = self.processor(image, input_boxes=input_boxes, return_tensors="pt", do_rescale=False, do_normalize=False).to(self.device)
+        return inputs, image.size
+
+    def inference(self, inputs):
+        with torch.no_grad():
+            outputs = self.model(**inputs, multimask_output=False)
+        return outputs
+
+    def postprocess(self, outputs, original_size):
+        # Post-process masks to original image size
+        masks = self.processor.post_process_masks(
+            outputs.pred_masks.cpu(),
+            torch.tensor([original_size]), # (W, H) -> (H, W)
+            outputs.reshaped_input_sizes.cpu()
+        )
+
+        # Convert masks to binary numpy arrays and then to base64 for JSON response
+        results = []
+        for mask_dict in masks:
+            mask_np = mask_dict['segmentation'].squeeze().numpy().astype(np.uint8) * 255 # Convert to 0/255
+            buffered = io.BytesIO()
+            Image.fromarray(mask_np).save(buffered, format="PNG")
+            encoded_mask = base64.b64encode(buffered.getvalue()).decode('utf-8')
+            results.append({"mask": encoded_mask, "score": mask_dict.get('score', 0.0)})
+        return json.dumps(results)
+
+
+# Example of how to use the handler locally (for testing)
+if __name__ == '__main__':
+    handler = InferenceHandler()
+
+    # Create a dummy image
+    dummy_image_size = (256, 256)
+    dummy_image_np = np.random.randint(0, 256, dummy_image_size, dtype=np.uint8)
+    image = Image.fromarray(dummy_image_np)
+
+    # Encode dummy image to base64
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    encoded_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+    # Example bounding box
+    example_boxes = [[50, 50, 200, 200]]
+
+    # Create a dummy request body
+    dummy_request_body = json.dumps({"image": encoded_image, "boxes": example_boxes})
+
+    print("
+--- Testing InferenceHandler locally ---")
+    inputs, original_size = handler.preprocess(dummy_request_body)
+    outputs = handler.inference(inputs)
+    processed_response = handler.postprocess(outputs, original_size)
+    print("Local test successful. Response structure (truncated):", processed_response[:200], "...")

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8604a4ffe7d3c99df24b1224a6c9593a4f1b82cd8b455eb64b6feb55103f498d
+oid sha256:2395d9f09a56238ae54dcb447b7c5230aedf8c7c46fff0644c28666901c6bc11
 size 374979376

requirements.txt

@@ -0,0 +1,8 @@
+
+torch==2.8.0+cpu
+transformers==4.57.1
+huggingface_hub==0.36.0
+nibabel==5.3.2
+numpy==2.0.2
+Pillow==10.3.0
+tqdm==4.67.1