leminhhung0101/shrimp_3stage

Shrimp Disease Detection and Lesion Segmentation --- 3-Stage Deep Learning Pipeline

A GPU-optimized three-stage deep learning pipeline for shrimp health classification, disease classification, and lesion segmentation from images.

Overview

The pipeline processes each image sequentially:

Input Image
    ↓
Stage 1 — EfficientNet-B0
Healthy / Disease
    ↓ if Disease
Stage 2 — ResNet18
BlackGill / WhiteSpot
    ↓
Stage 3 — MobileNetV3-Small U-Net
Binary Lesion Segmentation
    ↓
Mask / Polygon / Final Prediction

For each image, the system can produce:

• Final class: Healthy, BlackGill, or WhiteSpot
• Disease probability from Stage 1
• Disease-type probability from Stage 2
• Binary lesion segmentation mask from Stage 3
• Lesion polygons mapped back to the original image resolution
• Optional mask and overlay visualizations

---

Model Architecture

Stage 1 --- Healthy vs Disease

Backbone: EfficientNet-B0

Classifier head:

Dropout(0.35)
Linear(in_features, 2)

Classes:

0 = Healthy
1 = Disease

Main training configuration:

Input size       : 256 × 256
Batch size       : 128
Max epochs       : 15
Freeze epochs    : 2
Backbone LR      : 3e-5
Head LR          : 3e-4
Weight decay     : 5e-4
Label smoothing  : 0.05
Early stopping   : patience 3

The final decision threshold is selected using the validation set rather than being permanently fixed at 0.5.

---

Stage 2 --- BlackGill vs WhiteSpot

Stage 2 is applied only to images classified as Disease by Stage 1.

Backbone: ResNet18

Classifier head:

Dropout(0.35)
Linear(in_features, 2)

Classes:

0 = BlackGill
1 = WhiteSpot

Main training configuration:

Input size       : 256 × 256
Batch size       : 256
Max epochs       : 15
Freeze epochs    : 2
Backbone LR      : 3e-5
Head LR          : 3e-4
Weight decay     : 5e-4
Label smoothing  : 0.05
Early stopping   : patience 3

The WhiteSpot decision threshold is selected on the validation set.

---

Stage 3 --- Lesion Segmentation

Stage 3 uses a MobileNetV3-Small encoder with a U-Net-style decoder.

MobileNetV3-Small Encoder
        ↓
Multi-scale Skip Features
        ↓
U-Net Decoder
        ↓
1-Channel Lesion Mask

Model:

MobileNetV3SmallUNet

Output tensor:

[B, 1, 256, 256]

Main training configuration:

Input size       : 256 × 256
Batch size       : 24
Max epochs       : 20
Learning rate    : 2e-4
Weight decay     : 1e-4
Early stopping   : patience 4

The segmentation threshold is selected according to validation Dice score.

After inference, the predicted mask can be:

1. Thresholded into a binary mask.
2. Cleaned by removing very small connected components.
3. Converted to contours with OpenCV.
4. Approximated as polygons.
5. Rescaled to the original image dimensions.

---

FP16 NPY Data Pipeline

The training pipeline uses a shared FP16 NPY image store:

train_images_256_fp16.npy

Image format:

shape : [N, 3, 256, 256]
dtype : float16
range : [0, 1]
layout: CHW

JPEG/PNG images are decoded and resized only once during preprocessing.

During training, samples are read directly from the NPY store using memory mapping:

np.load(path, mmap_mode="r")

This avoids repeated JPEG decoding and resizing on every epoch.

Stage 3 uses a separate FP16 mask store:

stage3_masks_256_fp16.npy

Mask format:

shape : [N, 1, 256, 256]
dtype : float16
values: 0 / 1

The original segmentation annotations are parsed from coordinate strings in the form:

x1,y1,x2,y2,x3,y3,...

and rasterized into binary masks.

---

NVIDIA T4 Optimization

The notebook is designed to keep an NVIDIA T4 busy while data for upcoming batches is prepared.

Main optimizations include:

• Automatic Mixed Precision (AMP / FP16)
• Tensor Core utilization
• channels_last memory format
• Fused AdamW when supported
• Pinned host memory
• Persistent DataLoader workers
• DataLoader prefetching
• Non-blocking host-to-device transfers
• CUDA prefetch stream
• GPU-side augmentation
• GPU-side normalization
• NPY memory mapping
• Block-based shuffle to reduce random seeks on Google Drive
• Minimal CPU/GPU synchronization inside training loops
• Batch inference instead of per-image inference

Conceptual data flow:

Google Drive NPY
       ↓
DataLoader Workers
       ↓
Pinned RAM
       ↓
CUDA Prefetch Stream — Batch N+1
       ↓
NVIDIA T4 computes Batch N

Typical configuration:

NUM_WORKERS = 2
PREFETCH_FACTOR = 3
USE_CUDA_PREFETCH = True
USE_AMP = True
USE_CHANNELS_LAST = True
USE_FUSED_ADAMW = True

---

Full Retraining

After selecting the best epoch using the validation split:

Stage 1 → s1_best_epoch
Stage 2 → s2_best_epoch
Stage 3 → s3_best_epoch

each model is retrained using all eligible data for its stage.

Final model state files:

stage1_fulltrain_state.pt
stage2_fulltrain_state.pt
stage3_fulltrain_state.pt

---

Combined Checkpoint

The three stages can be packaged into one deployment checkpoint:

shrimp_pipeline.pt

A combined checkpoint may contain:

{
    "stage1_state_dict": ...,
    "stage2_state_dict": ...,
    "stage3_state_dict": ...,

    "stage1_threshold": ...,
    "stage2_threshold": ...,
    "mask_threshold": ...,

    "stage1_best_epoch": ...,
    "stage2_best_epoch": ...,
    "stage3_best_epoch": ...,

    "stage1_validation_macro_f1": ...,
    "stage2_validation_macro_f1": ...,
    "stage3_validation_dice": ...,

    "stage1_img_size": 256,
    "stage2_img_size": 256,
    "stage3_img_size": 256,

    "stage1_classes": ["Healthy", "Disease"],
    "stage2_classes": ["BlackGill", "WhiteSpot"],

    "parameter_count": {...},
}

This allows the complete inference pipeline to be distributed as a single model file.

---

Inference Logic

Simplified inference flow:

p_disease = stage1(image)

if p_disease < stage1_threshold:
    prediction = "Healthy"
    polygons = []

else:
    p_whitespot = stage2(image)

    if p_whitespot >= stage2_threshold:
        prediction = "WhiteSpot"
    else:
        prediction = "BlackGill"

    mask = sigmoid(stage3(image)) >= mask_threshold
    polygons = mask_to_polygons(mask)

Stage 2 and Stage 3 are skipped for images classified as Healthy, reducing inference cost.

---

Single-Image Inference

Example:

pipeline = ShrimpPipeline(
    checkpoint_path="shrimp_pipeline.pt",
    device=device,
)

result = pipeline.predict("sample.jpg")

print(result)

Example Disease output:

{
    "disease": "BlackGill",
    "p_disease": 0.9821,
    "p_whitespot": 0.1345,
    "polygons": [
        [[120, 80], [180, 75], [210, 130], [150, 160]]
    ],
}

Example Healthy output:

{
    "disease": "Healthy",
    "p_disease": 0.083,
    "p_whitespot": None,
    "polygons": [],
}

---

Folder / Batch Inference

For a folder containing many images, batch inference should be used instead of predicting one image at a time.

Recommended pattern:

All Images
    ↓ batch
Stage 1
    ↓
Disease Subset Only
    ↓ batch
Stage 2
    +
Stage 3

Typical inference batch sizes on an NVIDIA T4:

INFER_CLS_BATCH = 256
INFER_DISEASE_BATCH = 64

If CUDA runs out of memory, reduce the Stage 3 batch size first.

---

Prediction Outputs

A folder inference script can generate:

predictions/
├── predictions.csv
├── masks/
│   ├── image_001_mask.png
│   └── ...
└── overlays/
    ├── image_001_overlay.jpg
    └── ...

Example CSV fields:

file_name
disease_prob
stage1_label
blackgill_prob
whitespot_prob
stage2_label
lesion_ratio
prediction
mask_path
overlay_path

A submission-oriented output can also use:

image_id,disease,polygon

Example:

image_id,disease,polygon
001.jpg,Healthy,
002.jpg,BlackGill,"120,80,180,75,210,130,150,160"
003.jpg,WhiteSpot,"45,60,90,55,110,100,70,120"

---

Project Files

A typical project directory may look like:

shrimp_3stage/
│
├── shrimp_pipeline.pt
│
├── stage1_best.pt
├── stage2_best.pt
├── stage3_best.pt
│
├── stage1_fulltrain_state.pt
├── stage2_fulltrain_state.pt
├── stage3_fulltrain_state.pt
│
├── stage1_history.csv
├── stage2_history.csv
├── stage3_history.csv
│
├── train_ids_90.csv
├── valid_ids_10.csv
│
├── public_output.csv
├── private_output.csv
│
└── npy_fp16_256/
    ├── train_images_256_fp16.npy
    ├── train_images_256_index.csv
    ├── train_images_256_fp16.done.json
    ├── stage3_masks_256_fp16.npy
    ├── stage3_masks_256_index.csv
    ├── public_images_256_fp16.npy
    └── private_images_256_fp16.npy

---

Requirements

Recommended environment:

Python >= 3.10
PyTorch >= 2.x
torchvision
numpy
pandas
opencv-python
Pillow
scikit-learn
tqdm

Installation:

pip install torch torchvision numpy pandas opencv-python pillow scikit-learn tqdm

---

Hardware

The training notebook is primarily optimized for:

Google Colab
NVIDIA T4 — 16 GB VRAM

The models can also run on other CUDA GPUs or on CPU.

When CUDA is unavailable:

device = torch.device("cpu")

CUDA-specific optimizations such as AMP, Tensor Cores, asynchronous H2D prefetching, and CUDA streams will not be used.

---

Why FP16 NPY?

FP16 NPY is used to improve training throughput across repeated epochs.

Advantages:

• JPEG/PNG decoding is not repeated every epoch.
• Image resizing is not repeated every epoch.
• Storage is approximately half the size of equivalent float32 arrays.
• Data requires little CPU preprocessing before GPU transfer.
• Memory mapping avoids loading the entire dataset into RAM.

Trade-offs:

• NPY files are larger than compressed JPEG/PNG images.
• Performance can still be limited by Google Drive bandwidth.
• Large NPY files should be memory-mapped rather than fully loaded into RAM.

---

Validation Metrics

The pipeline tracks validation metrics separately for each task:

Stage 1: Macro-F1
Stage 2: Macro-F1
Stage 3: Dice score

The exact values depend on the training run and should be read from the generated checkpoints/history files rather than hard-coded into this README.

Example:

import torch

ckpt = torch.load(
    "shrimp_pipeline.pt",
    map_location="cpu",
)

print(
    "Stage 1 Macro-F1:",
    ckpt["stage1_validation_macro_f1"],
)

print(
    "Stage 2 Macro-F1:",
    ckpt["stage2_validation_macro_f1"],
)

print(
    "Stage 3 Dice:",
    ckpt["stage3_validation_dice"],
)

print(
    "Thresholds:",
    ckpt["stage1_threshold"],
    ckpt["stage2_threshold"],
    ckpt["mask_threshold"],
)

---

Deployment Recommendations

For high-throughput inference:

• Use batch inference.
• Keep AMP enabled on CUDA.
• Avoid loading models repeatedly.
• Load all three models once and reuse them.
• Run Stage 2 and Stage 3 only on Disease candidates.
• Use pinned memory and non-blocking GPU transfers.
• Use larger classification batches than segmentation batches.

Suggested NVIDIA T4 starting points:

Stage 1 batch: 128–256
Stage 2 batch: 128–256
Stage 3 batch: 24–64

---

Limitations

• The current pipeline supports three final classes:
  • Healthy
  • BlackGill
  • WhiteSpot
• Stage 3 quality depends strongly on the segmentation annotation quality.
• Validation-optimized thresholds may need recalibration when the image domain changes.
• Resizing all images to 256 × 256 may reduce sensitivity to very small lesions.
• Predictions should not be treated as a substitute for expert biological or veterinary assessment.

---

Pipeline Summary

                    ┌──────────────────┐
                    │   Input Image    │
                    └────────┬─────────┘
                             │
                             ▼
                  ┌─────────────────────┐
                  │ EfficientNet-B0     │
                  │ Stage 1             │
                  │ Healthy / Disease   │
                  └────────┬────────────┘
                           │
             Healthy ◄─────┴─────► Disease
                                  │
                                  ▼
                       ┌────────────────────┐
                       │ ResNet18           │
                       │ Stage 2            │
                       │ BlackGill /        │
                       │ WhiteSpot          │
                       └─────────┬──────────┘
                                 │
                                 ▼
                       ┌────────────────────┐
                       │ MobileNetV3-Small  │
                       │ U-Net              │
                       │ Stage 3            │
                       │ Lesion Mask        │
                       └─────────┬──────────┘
                                 │
                                 ▼
                      Mask → Polygon → Output

---

Recommended Deployment Model

For deployment, the recommended artifact is:

shrimp_pipeline.pt

It can contain all three model state dictionaries, image sizes, thresholds, validation metrics, class definitions, and other metadata required by the inference pipeline.