Baywatch — Marine Debris Detection

Train a lightweight ResNet18 multi-label classifier to detect marine debris (and cloud) from images. The repo includes:

• Data prep: parse Pascal VOC-style XML annotations into a CSV
• Training: train and save a model checkpoint + loss metrics
• Inference: run predictions on a folder of images
• Plotting: generate a training-loss curve

Quickstart

1) Install dependencies

python -m venv .venv
. .venv/bin/activate   # macOS/Linux
# OR on Windows PowerShell:
# .\.venv\Scripts\Activate.ps1

pip install -r requirements.txt

2) Prepare your data

Expected layout:

data/
  train/                # training images + XML annotations
    *.png|*.jpg|...      # images (any PIL-supported format)
    *.xml                # Pascal VOC-style annotation files
  test/                 # images to run inference on
    *.png                # inference.py currently scans for .png files

Generate the labels CSV from data/train/*.xml:

python clean_csv.py

This creates data/labels_parsed.csv with columns:

• filename: image filename relative to data/train
• debris: 0/1
• cloud: 0/1

3) Train

python main.py

Outputs:

• model.pth: saved model weights
• training_metrics.csv: epoch/loss log

Inference (predict on a folder)

Run predictions on data/test using the saved model.pth:

python inference.py

What you’ll see:

• Per-image scores for Debris and Cloud (sigmoid probabilities)
• An “ALERT” if debris score (\ge) the configured threshold

Plot training loss

After training:

python plot_metrics.py

Outputs:

• training_metrics.png: loss curve image (and also shows the plot interactively)

Configuration

These defaults live at the top of main.py / inference.py:

• Paths: data/train, data/test, data/labels_parsed.csv, model.pth
• Training: BATCH_SIZE=8, EPOCHS=12, lr=0.001
• Threshold: THRESHOLD=0.5
• Device: uses CUDA if available, otherwise CPU

Project files

main.py         # train (and writes metrics/model)
inference.py    # run predictions on data/test images
clean_csv.py    # build data/labels_parsed.csv from VOC XMLs
utils.py        # PyTorch Dataset (reads images + labels CSV)
plot_metrics.py # plot training_metrics.csv -> training_metrics.png
requirements.txt

Notes

• inference.py currently only considers files ending in .png. If your test images are .jpg, update the extension filter.
• Training uses BCEWithLogitsLoss with 2 outputs (debris, cloud) and applies sigmoid at inference time.