README.md

---
license: other
library_name: pytorch
pipeline_tag: image-segmentation
tags:
  - semantic-segmentation
  - visual-pollution
  - billboard
  - knowledge-distillation
  - sam3
  - shssnet
  - pytorch
---

# SHSSNet: SAM3-distilled visual-pollution (billboard) segmentation

Lightweight (7M to 16M parameter) binary segmentation networks, **background vs. visual pollution
(billboards and signage)**, trained with knowledge distillation from
[SAM 3](https://huggingface.co/facebook/sam3). SAM 3 is used only as a training-time teacher.
The students run standalone at inference and do not need SAM 3.

This repository holds two experiment folders:

| folder | what it is | files |
|---|---|---|
| `shssnet_distillation/` | Pretrain on `billboard.v1i.yolov11` (SAM3 pseudo-labels), fine-tunes on `dataset_khresna`, ablations, and a loss-weight grid | 23 `.pt` |
| `shssnet_kd_sam3/` | KD from SAM3 trained from scratch on `dataset_khresna` with human GT masks plus SAM3 feature/presence distillation: pilot runs, baselines and Kaggle runs | 9 `.pt` |

## Model variants

| variant | params | GFLOPs | notes |
|---|---|---|---|
| `shssnet_c` | 7.44M | 27.5 | compact |
| `shssnet_c_wide` | 13.03M | 40.2 | wider compact |
| `shssnet_p` | 15.68M | 68.9 | largest, best accuracy |

Run suffixes:
- `runA`: pseudo-labels from the *billboard-only* SAM3 prompt.
- `runB`: pseudo-labels from the full 7-concept SAM3 union.
- `ft_*`: fine-tuned on `dataset_khresna` (ds1 + ds2) at 512 px.
- `*_kd`: fine-tuned with distillation losses.
- `ablate_*`: single-component ablations of `shssnet_c_runA`. Each one disables a presence gate, a hard gate, the boundary skip, the feature loss, and so on.
- `grid_*_lf{λ_feature}_lb{λ_boundary}`: loss-weight grid on `shssnet_c_runB`.
- `pilot_*` and `pilot_base_*`: 20-epoch distilled and non-distilled baseline pairs.
- `kaggle_*`: runs trained on Kaggle (`_last` is the last epoch rather than the best).

Architecture: presence-gated, boundary-aware encoder-decoder with a shared stride-16 feature
(256 ch), a mask head, a boundary head and an image-level presence head.

## Distillation objective

```
L = λ_feature·L_feature + λ_presence·L_presence + λ_boundary·L_boundary + λ_seg·L_seg
```

- `L_feature`: cosine alignment of the student's shared feature to SAM3 FPN features.
- `L_presence`: student presence logit against the SAM3 image-level score.
- `L_boundary`: BCE+Dice on a morphological-gradient boundary band.
- `L_seg`: BCE+Dice on the mask.

## Results

### Held-out `seg_ds5_pakEvert` (967 field photos, never seen in training)

| checkpoint | fg IoU | Dice | mAP@50 (mask) | mAP@50-95 (mask) |
|---|---|---|---|---|
| `ft_p` | **0.669** | **0.802** | 0.324 | 0.218 |
| `shssnet_p_runB` | 0.655 | 0.792 | **0.393** | **0.257** |
| `shssnet_p_runA` | 0.609 | 0.757 | 0.308 | 0.204 |
| `ft_c_wide` | 0.499 | 0.666 | 0.178 | 0.095 |
| `ft_c` | 0.483 | 0.652 | 0.180 | 0.098 |
| `shssnet_c_runA` | 0.448 | 0.618 | 0.171 | 0.096 |

`ft_p` wins pixel metrics and `shssnet_p_runB` wins instance-level mAP. mAP is computed over
connected-component instances, since the task is binary, so treat it as indicative.
The 11 checkpoints evaluated here (all but the `shssnet_kd_sam3/` pilots) were never trained
on ds5 images. The `shssnet_kd_sam3/pilot_*` checkpoints did train on about 72% of ds5 and
**should not be compared on it**.

### Test split of the `billboard.v1i.yolov11` re-split (`shssnet_distillation/`)

| run | mIoU | Dice | params |
|---|---|---|---|
| `shssnet_p_runB` | 0.8785 | 0.9341 | 15.7M |
| `shssnet_p_runA` | 0.8536 | 0.9190 | 15.7M |
| `shssnet_c_runA` | 0.7959 | 0.8823 | 7.4M |
| `shssnet_c_wide_runA` | 0.7830 | 0.8735 | 13.0M |
| SegFormer-B1 baseline | 0.8868 | 0.9390 | n/a |
| PIDNet-S baseline | 0.7478 | 0.8495 | n/a |

Small objects remain the weak spot: recall@0.5 on small instances is about 0 to 0.11 even for the P variant.

## Usage

Checkpoints are **full training checkpoints** (`model_state`, `optimizer_state`, `scheduler_state`,
`epoch`, `metric_value`, ...). For `shssnet_kd_sam3/` they also include `loss_module_state`.
Weights are under `model_state`. Each file is roughly 90 MB to 190 MB because it carries optimizer state.

```python
import torch
from huggingface_hub import hf_hub_download

# repo_id is a placeholder: replace it with this repo's id
path = hf_hub_download("<user-or-org>/<repo>", "shssnet_distillation/ft_p.pt")
ckpt = torch.load(path, map_location="cpu", weights_only=False)
print(ckpt["epoch"], ckpt["metric_value"])

# model classes live in the `shssnet` package of the matching experiment folder
# from shssnet.models import build_model
# model = build_model(cfg); model.load_state_dict(ckpt["model_state"]); model.eval()
```

Use the `shssnet` package from the matching experiment (`shssnet_distillation` or
`shssnet_kd_sam3`) in the code repo: <github-repo-url>. The class-defining config is in that
experiment's `configs/`. `weights_only=False` is needed because the checkpoints contain RNG
state, so only load files you trust.

Input size: 640x640 for the pretrain runs, 512x512 for `ft_*` and `shssnet_kd_sam3` runs.
Output: a foreground mask logit `(B,1,H,W)`, a boundary logit `(B,1,H/4,W/4)` and an image-level presence logit `(B,)`.

## Training data

- **`billboard.v1i.yolov11`** (Roboflow): 1,721 human-reviewed images, re-split by augmentation group (~71/14/14).
  Pseudo-masks come from SAM3 (`facebook/sam3`) prompted per pollution sub-concept.
- **`dataset_khresna`** (`seg_ds1_mainBB`, `seg_ds2_indoRoboflow`, `seg_ds5_pakEvert`): 6,525 usable images after
  removing 64 messy ones, split 4,660/932/933 grouped by source photo.
- **`[Unsegment]Billboard_Dataset_pakEvert`**: 403 raw field photos used only for qualitative testing.

## Limitations

- Binary only (visual pollution vs. background). It does not classify the pollution type.
- Trained on billboard-style imagery from a few sources, so expect a domain gap elsewhere. IoU on field photos (0.67) is far below the in-domain test mIoU (0.88).
- Small objects are poorly recalled.
- Labels for the pretrain runs are SAM3 pseudo-labels and inherit the teacher's errors.
- SAM 3 is released under its own license; check `facebook/sam3` terms before commercial use. The `license: other` tag above is a placeholder, so set it to what you intend.

## Citation

```bibtex
@misc{shssnet_visual_pollution,
  title  = {SHSSNet: SAM3-distilled visual-pollution segmentation},
  author = {Rafi Abhista},
  year   = {2026},
  note   = {Model checkpoints}
}
```