ransoppong/halounet

Thyroid Ultrasound Segmentation Model

This repository contains a PyTorch-based U-Net model for segmenting thyroid nodules in ultrasound images. The model is designed to automatically identify and segment thyroid regions from ultrasound scans, providing both binary masks and confidence overlays.

Model Architecture

The model uses a U-Net architecture with the following specifications:

• Input: Grayscale ultrasound images (224x224 pixels)
• Output: Binary segmentation masks and confidence overlays
• Architecture: U-Net with skip connections
• Features: [64, 128, 256, 512] channels in the encoder/decoder paths
• Activation: ReLU with batch normalization
• Final Layer: 1x1 convolution with sigmoid activation

Model Details

• Framework: PyTorch
• Input Channels: 1 (grayscale)
• Output Channels: 1 (binary segmentation)
• Input Size: 224x224 pixels
• Model Size: ~31M parameters
• Task: Medical image segmentation (thyroid ultrasound)

Usage

Using the Model Directly

import torch
from PIL import Image
import numpy as np

# Load the model
model = torch.load('best_model.pth', map_location='cpu')
model.eval()

# Preprocess image
def preprocess_image(image_path):
    image = Image.open(image_path).convert('L')
    image = image.resize((224, 224))
    image_array = np.array(image, dtype=np.float32) / 255.0
    image_tensor = torch.from_numpy(image_array).unsqueeze(0).unsqueeze(0)
    return image_tensor

# Run inference
input_tensor = preprocess_image('your_image.png')
with torch.no_grad():
    prediction = model(input_tensor)
    mask = torch.sigmoid(prediction) > 0.5

Using the API

The model is also available as a FastAPI service with the following endpoints:

• POST /segment-image/ - Upload and segment an image file
• POST /segment-image-base64/ - Send base64 encoded image data
• POST /segment-image-batch/ - Process multiple images in batch
• GET /health - Health check endpoint

API Example

import requests

# Segment an image file
with open('thyroid_ultrasound.png', 'rb') as f:
    files = {'file': f}
    response = requests.post('http://localhost:8000/segment-image/', files=files)
    result = response.json()

# Get segmentation mask and overlay
mask = result['segmentation_mask']
overlay = result['overlay_image']
confidence = result['confidence_score']

Training

The model was trained on a dataset of thyroid ultrasound images with corresponding segmentation masks. The training process included:

• Data augmentation (rotation, scaling, flipping)
• Binary cross-entropy loss
• Adam optimizer
• Learning rate scheduling
• Early stopping based on validation performance

Performance

• Input Processing: Supports common image formats (PNG, JPG, BMP, TIFF)
• Output Quality: High-resolution segmentation masks
• Confidence Scoring: Provides confidence overlays with color coding
• Batch Processing: Supports processing multiple images simultaneously

Requirements

torch>=2.0.0
torchvision>=0.15.0
Pillow>=10.0.0
numpy>=1.24.0
fastapi>=0.104.1
uvicorn[standard]>=0.24.0

Citation

If you use this model in your research, please cite:

@misc{thyroid_segmentation_2024,
  title={Thyroid Ultrasound Segmentation using U-Net},
  author={Your Name},
  year={2024},
  url={https://huggingface.co/your-username/thyroid-segmentation}
}

License

This model is released under the MIT License. See the LICENSE file for details.

Contributing

Contributions are welcome! Please feel free to submit issues, feature requests, or pull requests.

Support

For questions or support, please open an issue on this repository or contact the maintainers.