M.Tech project (AI & Data Science, SASTRA Deemed University) exploring whether adding a graph-reasoning step at the bottleneck of a Swin-Transformer U-Net actually helps with underwater semantic segmentation, using the SUIM dataset as the benchmark.
Short version: yes, it helps — especially on the classes that are hardest to tell apart underwater (wrecks, robots/instruments, fish) — but not for free. There's a compute cost, and one class (aquatic plants/flora) stays stubbornly hard for every model we tried.
Most underwater segmentation work either bolts a CNN encoder onto a U-Net decoder or throws a plain ViT at the problem and calls it a day. Neither really models the fact that underwater scenes are relational — a diver next to a wreck, a fish near a reef, background water that spans the whole frame. The idea here was to keep a strong Swin encoder + U-Net decoder as the backbone, but add a lightweight GCN (EdgeConv over a k-NN graph built from the bottleneck features) so the model can reason about how regions relate to each other before decoding, rather than treating every spatial location independently.
That's TransGCN-UNet. Everything else in this repo — the DenseNet baseline, the two ablations — exists to answer "okay, but is the GCN actually doing anything?"
SUIM — 8 semantic classes:
| Short | Class |
|---|---|
| BG | Waterbody Background |
| HD | Human Divers |
| PF | Aquatic Plants & Flora |
| WR | Wrecks & Ruins |
| RO | Robots & Instruments |
| RI | Reefs & Invertebrates |
| FV | Fish & Vertebrates |
| SR | Sea-Floor & Rocks |
Images resized to 224×224 for the Swin-based models (TransGCN-UNet and both ablations); the DenseNet-201+UNet baseline uses its own preprocessing to match the original paper setup as closely as possible.
- DenseNet-201 + UNet — baseline encoder-decoder, no attention, no graph reasoning. Loss: 0.5×Dice + 0.5×Focal.
- TransGCN-UNet (full) — Swin-Base encoder → GCN refinement (3× EdgeConv,
k=9) at the bottleneck → U-Net decoder with deep supervision. Loss: 0.35×Dice
- 0.30×Focal + 0.35×Lovász.
- Swin+UNet, no GCN (ablation) — same encoder, same decoder, same loss — the GCN is swapped for a plain linear projection. Isolates what the graph reasoning is actually contributing.
- TransGCN-UNet, no Lovász (ablation) — full architecture, GCN included, but loss dropped to 0.5×Dice + 0.5×Focal to match the baseline's loss weighting. Isolates the loss function's contribution from the architecture's.
All four models are trained for the same number of epochs (60) with the same optimizer/LR schedule where applicable, so encoder, GCN, and loss effects can be read off independently instead of getting tangled together.
| Model | mIoU | Mean F1 | Pixel Acc |
|---|---|---|---|
| DenseNet-201 + UNet | 68.57% | 80.25% | 84.82% |
| TransGCN-UNet (full) | 72.41% | 82.61% | 86.94% |
| Swin+UNet (no GCN) | 73.07% | 83.15% | 88.20% |
| TransGCN-UNet (no Lovász) | see app/outputs/ |
Per-class numbers, ablation deltas, and the params/FLOPs/CPU-latency proxy
comparison are in the Streamlit dashboard (app/app.py). Aquatic Plants & Flora
is the weak point across every model (30–34% IoU); everything else clears 60%+.
Interestingly, Swin+UNet without the GCN slightly outperforms the full TransGCN-UNet on mIoU — worth digging into further rather than sweeping under the rug (see Known limitations below).
TRANSGCN_PROJECT/
└── app/
├── checkpoints/ # trained weights (Git LFS — see .gitattributes)
│ ├── DenseNet201_best.pth
│ ├── SwinUNet_noGCN_best.pth
│ ├── TransGCN_UNet_best.pth
│ └── TransGCN_UNet_noLovasz_best.pth
├── logs/ # per-epoch train/val history
│ ├── DenseNet201_history.json
│ ├── SwinUNet_noGCN_history.json
│ ├── TransGCN_UNet_history.json
│ └── TransGCN_UNet_noLovasz_history.json
├── outputs/ # test-set results consumed by app.py
│ ├── DenseNet201_combined_plots.png
│ ├── DenseNet201_test_results.json
│ ├── deployment_proxy_benchmark.json # FLOPs / params / CPU-latency proxy
│ ├── SwinUNet_noGCN_test_results.json
│ ├── TransGCN_UNet_combined_plots.png
│ ├── TransGCN_UNet_noLovasz_test_results.json
│ └── TransGCN_UNet_test_results.json
├── code_files/
│ ├── Project_DenseNet_TransGCN_UNet.ipynb # baseline + full TransGCN-UNet
│ └── TransGCN_Ablation_Study.ipynb # both ablations
├── app.py # Streamlit dashboard
├── dump_keys.py # inspects/validates test_results.json keys
├── .gitattributes
└── requirements.txt
cd app
pip install -r requirements.txt
streamlit run app.pyIt reads whatever *_test_results.json files it finds in outputs/ — if an
ablation variant's JSON isn't there yet, that tab just says so instead of
crashing. All JSON files use the same class-name key convention
(iou_Waterbody Background, iou_Aquatic Plants & Flora, etc.) so they're
interchangeable regardless of which notebook produced them.
dump_keys.py is a small sanity-check script — run it against any
outputs/*.json if a dashboard card looks wrong, to confirm the file actually
has the keys the dashboard expects before assuming it's a code bug.
Both notebooks are meant to run on Colab (they mount Drive for data + checkpoint storage, and expect a GPU runtime).
Project_DenseNet_TransGCN_UNet.ipynb— trains DenseNet-201+UNet and the full TransGCN-UNet, evaluates both on the test set.TransGCN_Ablation_Study.ipynb— trains (or, if you already have the.pthfiles, just evaluates) the two ablation variants and writes theirtest_results.json.
Training config (epochs, batch size, LR, optimizer) is matched across all four models on purpose — the point of the ablations is isolating one variable at a time, so if the schedules drifted apart the comparison wouldn't mean anything.
- The CPU-latency numbers in
deployment_proxy_benchmark.jsonare a compute proxy, not a measurement on real embedded hardware (e.g. a Jetson). Treat them as a rough upper bound, not a deployment claim. - Aquatic Plants & Flora is the consistent failure case — likely a mix of class imbalance and genuinely ambiguous boundaries in the source images (plants blending into background/rocks). Worth a dedicated look before calling this architecture "done."
- Swin+UNet (no GCN) currently edges out the full TransGCN-UNet on mIoU/F1/ PixAcc — the GCN's benefit shows up more clearly on individual hard classes (wrecks, robots, fish) than in the aggregate mIoU. Worth reporting honestly rather than cherry-picking the metric that favors the full model.
- No public release of the trained checkpoints is planned at this point (see the code availability note in the manuscript); this repo is for reproducibility/reference, not a ready-to-deploy package.
Manuscript based on this work is currently under revision at The Visual Computer, addressing reviewer comments from the initial suitability assessment. The two ablations in this repo exist specifically to answer reviewer questions about the GCN's and the Lovász loss's individual contributions.
MIT — see LICENSE. Trained checkpoints are not covered by this license and are not being publicly released (see above).
Vitesh — M.Tech, Artificial Intelligence & Data Science, SASTRA Deemed University. Advised by Dr. N. Sasikaladevi and Dr. Kannan Balasubramanian.