|
|
--- |
|
|
license: mit |
|
|
language: en |
|
|
library_name: transformers |
|
|
tags: |
|
|
- unet |
|
|
- film |
|
|
- computer-vision |
|
|
- image-segmentation |
|
|
- medical-imaging |
|
|
- pytorch |
|
|
--- |
|
|
|
|
|
# FILMUnet2D |
|
|
|
|
|
This model is a 2D U-Net with FiLM conditioning for Ultrasound multi-organ segmentation. |
|
|
|
|
|
## Installation |
|
|
|
|
|
Make sure you have `transformers` and `torch` installed. |
|
|
|
|
|
```bash |
|
|
pip install transformers torch |
|
|
``` |
|
|
|
|
|
## Usage |
|
|
|
|
|
You can load the model and processor using the `Auto` classes from `transformers`. Since this repository contains custom code, make sure to pass `trust_remote_code=True`. |
|
|
|
|
|
```python |
|
|
import torch |
|
|
from transformers import AutoModel, AutoImageProcessor |
|
|
from PIL import Image |
|
|
|
|
|
# 1. Load model and processor |
|
|
repo_id = "AImageLab-Zip/US_FiLMUNet" |
|
|
|
|
|
processor = AutoImageProcessor.from_pretrained(repo_id, trust_remote_code=True) |
|
|
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True) |
|
|
model.eval() |
|
|
|
|
|
# 2. Load and preprocess your image |
|
|
# The processor handles resizing, letterboxing, and normalization. |
|
|
image = Image.open("path/to/your/image.png").convert("RGB") |
|
|
inputs = processor(images=image, return_tensors="pt") |
|
|
|
|
|
# 3. Prepare conditioning input |
|
|
# This should be an integer tensor representing the organ ID. |
|
|
# Replace `4` with the appropriate ID for your use case. |
|
|
organ_id = torch.tensor([4]) |
|
|
|
|
|
# 4. Run inference |
|
|
with torch.no_grad(): |
|
|
outputs = model(**inputs, organ_id=organ_id) |
|
|
|
|
|
# 5. Post-process the output to get the final segmentation mask |
|
|
# The processor can convert the logits to a binary mask, automatically handling |
|
|
# the removal of letterbox padding and resizing to the original image dimensions. |
|
|
mask = processor.post_process_semantic_segmentation( |
|
|
outputs, |
|
|
inputs, |
|
|
threshold=0.7, |
|
|
return_as_pil=True |
|
|
)[0] |
|
|
|
|
|
# 6. Save the result |
|
|
mask.save("output_mask.png") |
|
|
|
|
|
print("Segmentation mask saved to output_mask.png") |
|
|
``` |
|
|
|
|
|
### Model Details |
|
|
|
|
|
- **Architecture:** U-Net with FiLM layers for conditional segmentation. |
|
|
- **Conditioning:** The model's output is conditioned on an `organ_id` input. |
|
|
- **Input:** RGB images. |
|
|
- **Output:** A single-channel segmentation mask. |
|
|
|
|
|
### Configuration |
|
|
|
|
|
The model configuration can be accessed via `model.config`. Key parameters include: |
|
|
- `in_channels`: Number of input channels (default: 3). |
|
|
- `num_classes`: Number of output classes (default: 1). |
|
|
- `n_organs`: The number of different organs the model was trained to condition on. |
|
|
- `depth`: The depth of the U-Net. |
|
|
- `size`: The base number of filters in the first layer. |
|
|
|
|
|
### Organ IDs |
|
|
|
|
|
The `organ_id` passed to the model corresponds to the following mapping: |
|
|
|
|
|
```python |
|
|
organ_to_class_dict = { |
|
|
"appendix": 0, |
|
|
"breast": 1, |
|
|
"breast_luminal": 1, |
|
|
"cardiac": 2, |
|
|
"thyroid": 3, |
|
|
"fetal": 4, |
|
|
"kidney": 5, |
|
|
"liver": 6, |
|
|
"testicle": 7, |
|
|
} |
|
|
``` |
|
|
|
|
|
### Alternative Versions |
|
|
|
|
|
This repository contains multiple versions of the model located in subfolders. You can load a specific version by using the `subfolder` parameter. |
|
|
|
|
|
#### 4-Stage U-Net |
|
|
|
|
|
This version has a U-Net depth of 4. |
|
|
|
|
|
```python |
|
|
from transformers import AutoModel |
|
|
|
|
|
model_4_stages = AutoModel.from_pretrained( |
|
|
"AImageLab-Zip/US_FiLMUNet", |
|
|
subfolder="unet_4_stages", |
|
|
trust_remote_code=True |
|
|
) |
|
|
``` |
|
|
|
