FISH LICE DETECTION

A technical tour of a two-stage salmon-lice vision pipeline

github.com/AtteberryShane/FishLiceDetection

The 30-second version

This repository builds a practical two-stage computer-vision system for Atlantic salmon lice. Instead of asking one model to both recognize an infected fish and pinpoint the problem, it divides the job. Stage 1 is a whole-image EfficientNet-B0 classifier: fresh or infected? If the image clears an infection-probability gate, Stage 2 runs another EfficientNet-B0 over overlapping image patches. High-scoring patches become candidate lice regions, and non-maximum suppression (NMS) removes redundant overlapping boxes.

Why the design works: the first stage is a cheap screening step; the second spends compute only where it matters. It is simple, inspectable, and easy to tune—but it is technically a patch-classification detector rather than a learned bounding-box detector.

Snapshot

Backbone

EfficientNet-B0 for both stages

Framework

PyTorch + torchvision + OpenCV

Stage 1 result

95.36% test accuracy in saved metrics

Stage 2 result

95.45% best validation accuracy in saved metrics

Interface

CLI + Gradio web GUI

1. Architecture: screen first, zoom in second

The repository’s core idea is deliberately modular. A whole-fish classifier handles the broad question, while a patch classifier handles localization. That split keeps the code compact and lets threshold tuning happen without retraining.

Stage 1 — “Does this fish look infected?”

`train_classifier.py` loads three SalmonScan-derived splits: `SalmonScan_train`, `SalmonScan_unseen` (used as validation), and `SalmonScan_test`. Images are resized/cropped to 224×224, normalized with ImageNet statistics, and augmented with flips, rotation, color jitter, and translation.

model = build_efficientnet_model(
    num_classes=len(train_dataset.classes),
    pretrained=args.use_pretrained,
    freeze_backbone=args.freeze_backbone,
)
criterion = nn.CrossEntropyLoss()
optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

View train_classifier.py  •  View models.py

Stage 2 — “Where are the suspicious patches?”

The second model is trained on `LicePositive` and `LiceNegative` image patches. At inference time it is applied repeatedly across a larger fish image. The scanning loop converts classification probabilities into coarse fixed-size bounding boxes.

for (x1, y1, x2, y2, crop) in windows:
    patch_tensor = preprocess(crop, args.detector_image_size).to(device)
    det_logits = detector(patch_tensor)
    lice_prob = float(F.softmax(det_logits, dim=1)[0, lice_idx])

    if lice_prob >= args.detection_threshold:
        boxes.append([x1, y1, x2 - x1, y2 - y1])
        scores.append(lice_prob)

View two_stage_infer.py  •  View train_detector.py

NMS cleans up the pile-up

Overlapping windows can all fire on the same lesion or lice region. OpenCV’s NMSBoxes keeps the strongest candidates while suppressing boxes that overlap too heavily. That is why the final images show a smaller set of red rectangles instead of a wall of nearly identical windows.

indices = cv2.dnn.NMSBoxes(
    boxes, scores,
    score_threshold=0.0,
    nms_threshold=threshold,
)

Translation: Stage 2 is not predicting box coordinates. The sliding window defines the box; the network only decides whether each window looks lice-positive.

The two stages live on very different data scales. Stage 1 has enough examples to support a conventional train/validation/test workflow. Stage 2 is tiny: 89 negative patches and just 21 positive patches. The code compensates with aggressive augmentation—vertical and horizontal flips, rotation, color jitter, translation, scale changes, and random perspective.

The shared model factory

def build_efficientnet_model(num_classes, pretrained=True,
                            freeze_backbone=False, dropout=0.5):
    weights = EfficientNet_B0_Weights.DEFAULT if pretrained else None
    model = efficientnet_b0(weights=weights)
    model.classifier[1] = nn.Linear(model.classifier[1].in_features,
                                    num_classes)
    model.classifier[0] = nn.Dropout(dropout)
    return model

This is a clean transfer-learning pattern: keep EfficientNet-B0’s feature extractor, replace the final classification layer, and optionally freeze the backbone. Dropout is increased to 0.5, explicitly noted in the source as an anti-overfitting choice.

Training loop

`engine.py` is intentionally generic. It runs the same basic training/evaluation loop for both stages, tracks validation accuracy, retains the best model state, and feeds validation loss to `ReduceLROnPlateau`. Checkpoints store model weights, class mappings, arguments, and metrics—good ingredients for reproducible inference.

One thing to watch: Stage 2 uses a seeded random 80/20 split, but it is not stratified. With only 21 positive patches, the exact class mix in validation matters a lot. There is also no independent Stage 2 test set in the current training script.

Git LFS note

Most dataset PNGs in the downloaded ZIP are Git LFS pointer files rather than the full image bytes. The generated `outputs/` PNGs are present normally. Anyone cloning the project for training should make sure Git LFS downloads the underlying datasets.

Repository README / setup commands

Stage 1: strong held-out performance

Best validation accuracy: 99.0%.  Test accuracy: 95.36%. Training accuracy rises quickly from ~75.8% to the high 90s, while validation stays consistently strong after the first few epochs. The gap between the best validation score and final test result is a useful reminder that a single validation peak is not the same thing as final generalization.

Stage 2: high validation accuracy on a very small split

Best validation accuracy: 95.45%. The patch model reaches 100% training accuracy at epoch 9 and ~95.5% validation accuracy at epochs 9–10. That looks encouraging, but the validation set is only about 22 patches. One misclassified sample therefore moves accuracy by roughly 4.5 percentage points. With such a small positive class, confusion matrices, precision/recall, and repeated or stratified splits would tell a much richer story than accuracy alone.

Metric reality check: For a lice detector, false negatives and false positives are operationally different. Accuracy collapses both into one number. The next evaluation pass should emphasize recall, precision, F1, PR curves, and image-level detection metrics.

What the saved output looks like

Figure 3. Example inference outputs from the repository. Red rectangles are retained sliding-window detections after thresholding and NMS.

4. Inference & UX: a detector you can actually poke at

The CLI is useful for batch experiments; the Gradio app makes threshold behavior visible. The GUI exposes the knobs that matter most: Stage 1 threshold, Stage 2 threshold, window size, stride, NMS threshold, positive class labels, and even a Stage 1 bypass for partial-fish images.

Inference controls

Control

Effect

Trade-off

Classification threshold

Gate into Stage 2

Higher = fewer images scanned

Detection threshold

Patch must exceed this score

Higher = fewer false positives, lower recall

Window size

Physical region each patch covers

Large = more context; small = finer localization

Stride

Distance between adjacent windows

Small = denser scan, slower inference

NMS threshold

Overlap suppression strength

Tunes duplicate-box removal

The GUI goes beyond “box or no box”

`gui_app.py` computes the maximum, mean, median, standard deviation, 95th percentile, and 99th percentile of Stage 2 window probabilities. It also groups windows into confidence bands and warns when the detector produces many candidate boxes. That is a genuinely useful debugging layer: it helps distinguish “the threshold is wrong” from “the model thinks everything is a louse.”

View gui_app.py

Run it

python src/two_stage_infer.py \
  --stage1-checkpoint src/checkpoints/stage1_classifier.pt \
  --stage2-checkpoint src/checkpoints/stage2_detector.pt \
  --input FishLiceDataSet/SalmonScan_unseen/InfectedFish \
  --output-dir outputs \
  --classification-threshold 0.55 \
  --detection-threshold 0.7

python src/gui_app.py \
  --stage1-checkpoint src/checkpoints/stage1_classifier.pt \
  --stage2-checkpoint src/checkpoints/stage2_detector.pt \
  --device cpu

Performance note

The current implementation scores sliding windows one at a time. That keeps the code obvious but leaves speed on the table. Batching dozens of patches into a single tensor would reduce Python/model-call overhead substantially, especially on GPU.

5. Engineering review: what is good, what I would change next

What is already solid

• Clear stage boundaries. Data loading, model construction, training, and inference are separated into focused modules.

• Practical transfer learning. EfficientNet-B0 is lightweight enough for a prototype while still offering strong pretrained image features.

• Reproducible artifacts. Checkpoints retain class mappings, CLI arguments, and metrics instead of saving naked weights only.

• Real debugging affordances. The GUI exposes probability distributions and threshold tuning rather than hiding everything behind one “detect” button.

• Edge-aware scanning. The sliding-window code adds bottom/right edge windows so the scan does not silently miss image borders.

Highest-value next steps

1. Make Stage 2 evaluation harder to fool — Use a stratified split or cross-validation, reserve a true test set, and report precision/recall/F1 plus confusion matrices. With 21 positive patches, each positive example is precious.

2. Address class imbalance explicitly — Try weighted cross-entropy or a weighted sampler; also grow the positive patch set. Augmentation creates variety, but it does not create genuinely new biological examples.

3. Batch sliding-window inference — Preprocess many windows, stack them, and score a batch in one forward pass. This is the easiest route to a large speedup without changing the model.

4. Upgrade localization when the dataset allows it — If bounding-box annotations can be created, compare against a true object detector such as YOLO, Faster R-CNN, RetinaNet, or a segmentation approach. The current boxes inherit the window geometry.

5. Calibrate thresholds on a target operating point — Choose thresholds from validation precision/recall behavior rather than from intuition alone. A farm monitoring use case may prefer very high recall; a manual-review workflow may tolerate more false positives.

6. Add experiment metadata — Record seed, commit hash, dataset version, training hardware, class counts, and chosen thresholds with each checkpoint. The current checkpoint structure is a good starting point.

Bottom line: This is a strong prototype architecture: understandable, runnable, and easy to interrogate. The biggest gap is not the neural network—it is evaluation depth and Stage 2 data volume. A larger, carefully held-out positive patch set would do more for confidence in the system than another round of model complexity.

Repository map

README

Data pipeline

Training engine

Model factory

Stage 1 training

Stage 2 training

Two-stage inference

Gradio interface