Architecture

A modular, secure design built for medical image segmentation with production deployment in mind.

U-Net Model Architecture

Encoder-Decoder with Skip Connections

  • 1Input: RGB image [3, 256, 256]
  • 2Encoder: 3-level downsampling with max pooling
  • 3Bottleneck: Deep feature extraction at lowest resolution
  • 4Decoder: 3-level upsampling with skip connections
  • 5Output: Logits [1, 256, 256]
  • 6Loss: BCE + 0.5 × Dice
Input [3,256,256]
64 ch
skip
128 ch
skip
256 ch
skip
512 ch
256 ch
128 ch
64 ch
Mask [1,256,256]

Module Structure

model.py

U-Net architecture

data.py

Dataset loading and augmentation

train.py

Training and finetuning loops

losses.py

Dice loss implementation

policy.py

3-state decision logic and policy locking

eval.py

Threshold calibration and test bucketing

predict.py

Single-image and batch inference

gallery.py

TP/FN/FP/TN bucket visualization grids

export.py

ONNX export with sigmoid wrapper and metadata

stats.py

Wound metrics (area, perimeter, circularity, color)

viz.py

Visualization utilities

config.py

Path configuration and dataclasses

io.py

Secure file I/O, path validation, checkpoint loading

validation.py

Input validation and sanitization

constants.py

Application-wide constants and limits

log.py

Centralized logging configuration

cli.py

Command-line interface

Security Features

Safe Model Loading

Uses weights_only=True (PyTorch 2.6+) to prevent arbitrary code execution from malicious checkpoint files.

Path Traversal Protection

All file paths are validated through validate_path() to prevent directory traversal attacks.

Input Validation

All user inputs (thresholds, sizes, learning rates, etc.) are validated with proper bounds checking.

Constants Module

Centralized constants with predefined limits (e.g., MAX_IMAGE_SIZE, MAX_EPOCHS).

Data Format

CSV Dataset Format

new_id

image_001

image_002

image_003

...

CSV files with new_id column mapping to image/mask filenames.

Directory Structure

project/

images/

image_001.png

image_002.png

masks/

image_001.png

image_002.png

dataset.csv