colosseum-microcontroller 0.1.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. colosseum_microcontroller-0.1.3/LICENSE +21 -0
  2. colosseum_microcontroller-0.1.3/PKG-INFO +54 -0
  3. colosseum_microcontroller-0.1.3/PYPI_README.md +36 -0
  4. colosseum_microcontroller-0.1.3/README.md +135 -0
  5. colosseum_microcontroller-0.1.3/akros/__init__.py +5 -0
  6. colosseum_microcontroller-0.1.3/akros/api.py +7 -0
  7. colosseum_microcontroller-0.1.3/colosseum/__init__.py +5 -0
  8. colosseum_microcontroller-0.1.3/colosseum/__main__.py +7 -0
  9. colosseum_microcontroller-0.1.3/colosseum/_core/__init__.py +1 -0
  10. colosseum_microcontroller-0.1.3/colosseum/_core/mergdet.py +27 -0
  11. colosseum_microcontroller-0.1.3/colosseum/_core/mergdet_m.py +191 -0
  12. colosseum_microcontroller-0.1.3/colosseum/_core/plate_dataset.py +354 -0
  13. colosseum_microcontroller-0.1.3/colosseum/_core/plate_loss.py +647 -0
  14. colosseum_microcontroller-0.1.3/colosseum/_core/plate_metrics.py +182 -0
  15. colosseum_microcontroller-0.1.3/colosseum/_core/platedet_u.py +429 -0
  16. colosseum_microcontroller-0.1.3/colosseum/_core/train_platedet.py +538 -0
  17. colosseum_microcontroller-0.1.3/colosseum/_core/training_state.py +96 -0
  18. colosseum_microcontroller-0.1.3/colosseum/api.py +159 -0
  19. colosseum_microcontroller-0.1.3/colosseum/cli.py +22 -0
  20. colosseum_microcontroller-0.1.3/colosseum/evaluation.py +68 -0
  21. colosseum_microcontroller-0.1.3/colosseum/pack_checkpoint.py +50 -0
  22. colosseum_microcontroller-0.1.3/colosseum/prediction.py +54 -0
  23. colosseum_microcontroller-0.1.3/colosseum/training.py +170 -0
  24. colosseum_microcontroller-0.1.3/colosseum/weights/amlpr.pt +0 -0
  25. colosseum_microcontroller-0.1.3/colosseum/weights/ccpd.pt +0 -0
  26. colosseum_microcontroller-0.1.3/colosseum/weights/drone.pt +0 -0
  27. colosseum_microcontroller-0.1.3/colosseum/weights/gtsdb.pt +0 -0
  28. colosseum_microcontroller-0.1.3/colosseum/weights/manifest.json +152 -0
  29. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/amlpr_seed1.pt +0 -0
  30. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/amlpr_seed2.pt +0 -0
  31. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/ccpd_seed1.pt +0 -0
  32. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/ccpd_seed2.pt +0 -0
  33. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/drone_seed1.pt +0 -0
  34. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/drone_seed2.pt +0 -0
  35. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/gtsdb_seed1.pt +0 -0
  36. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/gtsdb_seed2.pt +0 -0
  37. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/shwd_seed1.pt +0 -0
  38. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/shwd_seed2.pt +0 -0
  39. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/uc3m_seed1.pt +0 -0
  40. colosseum_microcontroller-0.1.3/colosseum/weights/seeds/uc3m_seed2.pt +0 -0
  41. colosseum_microcontroller-0.1.3/colosseum/weights/shwd.pt +0 -0
  42. colosseum_microcontroller-0.1.3/colosseum/weights/uc3m.pt +0 -0
  43. colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/PKG-INFO +54 -0
  44. colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/SOURCES.txt +48 -0
  45. colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/dependency_links.txt +1 -0
  46. colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/entry_points.txt +2 -0
  47. colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/requires.txt +3 -0
  48. colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/top_level.txt +2 -0
  49. colosseum_microcontroller-0.1.3/pyproject.toml +33 -0
  50. colosseum_microcontroller-0.1.3/setup.cfg +4 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Aqshal Nur Ikhsan
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,54 @@
1
+ Metadata-Version: 2.4
2
+ Name: colosseum-microcontroller
3
+ Version: 0.1.3
4
+ Summary: Train compact single-class object detectors for microcontrollers
5
+ Author: Aqshal Nur Ikhsan
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/Aqshalikhsan/Colosseum-Microcontroller
8
+ Project-URL: Repository, https://github.com/Aqshalikhsan/Colosseum-Microcontroller
9
+ Project-URL: Documentation, https://github.com/Aqshalikhsan/Colosseum-Microcontroller/blob/main/README.md
10
+ Keywords: object detection,microcontroller,ESP32-S3,TinyML,PyTorch
11
+ Requires-Python: >=3.11
12
+ Description-Content-Type: text/markdown
13
+ License-File: LICENSE
14
+ Requires-Dist: torch>=2.1
15
+ Requires-Dist: Pillow>=10
16
+ Requires-Dist: PyYAML>=6
17
+ Dynamic: license-file
18
+
19
+ <p align="center">
20
+ <img src="https://raw.githubusercontent.com/Aqshalikhsan/Colosseum-Microcontroller/main/assets/colosseum_banner.png" width="900" alt="Colosseum letter-tile logo">
21
+ </p>
22
+
23
+ Colosseum is a compact single-class object detection framework designed for microcontroller deployment. Install it to train a detector for your own target object. The package also includes 18 example checkpoints from six datasets and three seeds.
24
+
25
+ ```bash
26
+ python -m pip install colosseum-microcontroller
27
+ ```
28
+
29
+ ```python
30
+ from colosseum import Colosseum
31
+
32
+ model = Colosseum()
33
+ checkpoints = model.train(
34
+ data="my_dataset/data.yaml",
35
+ name="my_object",
36
+ label="my object",
37
+ grid="4x4",
38
+ )
39
+ print(checkpoints[0])
40
+ print(model.metrics) # AP50, AP75, AP50:95, precision, recall, F1
41
+ boxes = model.predict("path/to/image.jpg")
42
+ ```
43
+
44
+ Training starts from scratch, saves a compact checkpoint and `metrics.json` under `runs/custom/my_object/seed0/`, and loads it for inference. Metrics are computed on the validation split. One checkpoint detects one target class. Your dataset uses images and normalized bounding-box text annotations; see the [training guide](https://github.com/Aqshalikhsan/Colosseum-Microcontroller/blob/main/docs/TRAINING.md) for the layout.
45
+
46
+ The command-line equivalent is:
47
+
48
+ ```bash
49
+ colosseum train --data my_dataset/data.yaml --name my_object --label "my object" --grid 4x4
50
+ ```
51
+
52
+ The bundled example names are `uc3m`, `gtsdb`, `shwd`, `ccpd`, `drone`, and `amlpr`, plus `_seed1` and `_seed2` variants. For inference you can pass either one of these names or the path to your own compact checkpoint. See the [repository](https://github.com/Aqshalikhsan/Colosseum-Microcontroller) for benchmarks. The weights are floating-point PyTorch checkpoints; ESP32-S3 deployment uses the separate INT8 runtime.
53
+
54
+ [Open the training notebook in Google Colab](https://colab.research.google.com/github/Aqshalikhsan/Colosseum-Microcontroller/blob/main/notebooks/Train_Colosseum.ipynb) to upload a dataset and run the full pipeline.
@@ -0,0 +1,36 @@
1
+ <p align="center">
2
+ <img src="https://raw.githubusercontent.com/Aqshalikhsan/Colosseum-Microcontroller/main/assets/colosseum_banner.png" width="900" alt="Colosseum letter-tile logo">
3
+ </p>
4
+
5
+ Colosseum is a compact single-class object detection framework designed for microcontroller deployment. Install it to train a detector for your own target object. The package also includes 18 example checkpoints from six datasets and three seeds.
6
+
7
+ ```bash
8
+ python -m pip install colosseum-microcontroller
9
+ ```
10
+
11
+ ```python
12
+ from colosseum import Colosseum
13
+
14
+ model = Colosseum()
15
+ checkpoints = model.train(
16
+ data="my_dataset/data.yaml",
17
+ name="my_object",
18
+ label="my object",
19
+ grid="4x4",
20
+ )
21
+ print(checkpoints[0])
22
+ print(model.metrics) # AP50, AP75, AP50:95, precision, recall, F1
23
+ boxes = model.predict("path/to/image.jpg")
24
+ ```
25
+
26
+ Training starts from scratch, saves a compact checkpoint and `metrics.json` under `runs/custom/my_object/seed0/`, and loads it for inference. Metrics are computed on the validation split. One checkpoint detects one target class. Your dataset uses images and normalized bounding-box text annotations; see the [training guide](https://github.com/Aqshalikhsan/Colosseum-Microcontroller/blob/main/docs/TRAINING.md) for the layout.
27
+
28
+ The command-line equivalent is:
29
+
30
+ ```bash
31
+ colosseum train --data my_dataset/data.yaml --name my_object --label "my object" --grid 4x4
32
+ ```
33
+
34
+ The bundled example names are `uc3m`, `gtsdb`, `shwd`, `ccpd`, `drone`, and `amlpr`, plus `_seed1` and `_seed2` variants. For inference you can pass either one of these names or the path to your own compact checkpoint. See the [repository](https://github.com/Aqshalikhsan/Colosseum-Microcontroller) for benchmarks. The weights are floating-point PyTorch checkpoints; ESP32-S3 deployment uses the separate INT8 runtime.
35
+
36
+ [Open the training notebook in Google Colab](https://colab.research.google.com/github/Aqshalikhsan/Colosseum-Microcontroller/blob/main/notebooks/Train_Colosseum.ipynb) to upload a dataset and run the full pipeline.
@@ -0,0 +1,135 @@
1
+ <p align="center">
2
+ <img src="assets/colosseum_banner.png" width="900" alt="Colosseum letter-tile logo">
3
+ </p>
4
+
5
+ <p align="center">
6
+ Compact object detection for microcontrollers, with reproducible training and benchmarks.
7
+ </p>
8
+
9
+ <p align="center">
10
+ <a href="#quick-start"><img alt="Python 3.11" src="https://img.shields.io/badge/Python-3.11-3776AB?logo=python&amp;logoColor=white"></a>
11
+ <a href="#model-zoo"><img alt="PyTorch models" src="https://img.shields.io/badge/PyTorch-models-EE4C2C?logo=pytorch&amp;logoColor=white"></a>
12
+ <a href="#benchmarks"><img alt="ESP32-S3 benchmark" src="https://img.shields.io/badge/ESP32--S3-benchmark-00979D?logo=espressif&amp;logoColor=white"></a>
13
+ <a href="LICENSE"><img alt="MIT license" src="https://img.shields.io/badge/License-MIT-2DA44E"></a>
14
+ </p>
15
+
16
+ <p align="center">
17
+ <a href="#quick-start">🚀 Quick start</a> &nbsp;·&nbsp;
18
+ <a href="notebooks/Train_Colosseum.ipynb">📓 Training notebook</a> &nbsp;·&nbsp;
19
+ <a href="#model-zoo">🧠 Models</a> &nbsp;·&nbsp;
20
+ <a href="#benchmarks">📊 Benchmarks</a> &nbsp;·&nbsp;
21
+ <a href="#train-on-your-data">🛠️ Train on your data</a>
22
+ </p>
23
+
24
+ **Colosseum** is a compact, single-class object detector designed for microcontroller deployment. It uses fixed image-gradient features, residual convolution blocks, three-scale fusion, and separate confidence and box heads. The network has **81,403 parameters**. A checkpoint detects one target class; use a different checkpoint for a different class.
25
+
26
+ Install Colosseum as a Python library to train a detector on your own images. The repository also provides **18 ready-to-use PyTorch checkpoints** (six datasets × three seeds) for examples and benchmarks. Input is grayscale, letterboxed to **256 × 192** pixels. Colosseum outputs bounding boxes, not text recognition.
27
+
28
+ ## Quick start
29
+
30
+ Tested with Python 3.11. Install the package from PyPI:
31
+
32
+ ```bash
33
+ python -m pip install colosseum-microcontroller
34
+ ```
35
+
36
+ [Open the training notebook in Google Colab](https://colab.research.google.com/github/Aqshalikhsan/Colosseum-Microcontroller/blob/main/notebooks/Train_Colosseum.ipynb) to upload your own dataset, train through the Python API, view validation metrics, and download the checkpoints.
37
+
38
+ Create `my_dataset/data.yaml` and train your own detector from Python:
39
+
40
+ ```python
41
+ from colosseum import Colosseum
42
+
43
+ model = Colosseum()
44
+ checkpoints = model.train(
45
+ data="my_dataset/data.yaml",
46
+ name="my_object",
47
+ label="my object",
48
+ grid="4x4",
49
+ )
50
+ print(checkpoints[0]) # runs/custom/my_object/seed0/colosseum.pt
51
+ print(model.metrics) # AP50, AP75, AP50:95, precision, recall, F1
52
+ boxes = model.predict("path/to/image.jpg")
53
+ ```
54
+
55
+ This trains the Colosseum network from scratch, evaluates the validation split, loads the resulting checkpoint, and returns its path. Each seed writes a `metrics.json` with AP50, AP75, AP50:95, precision, recall, and F1. Call `model.evaluate("my_dataset/data.yaml")` to evaluate a loaded checkpoint again. Each checkpoint detects **one target class**. The dataset uses ordinary image files and text files with normalized box coordinates; the [training guide](docs/TRAINING.md) shows the exact layout. No Ultralytics package or YOLO model is used. For a source checkout, use `python -m pip install -e .` to install the library in editable mode.
56
+
57
+ The equivalent command-line interface is:
58
+
59
+ ```bash
60
+ colosseum train --data my_dataset/data.yaml --name my_object --label "my object" --grid 4x4
61
+ colosseum predict --weights runs/custom/my_object/seed0/colosseum.pt --source path/to/image.jpg
62
+ ```
63
+
64
+ The pip package also includes all **18 example checkpoints**, so a model name such as `uc3m`, `shwd`, or `uc3m_seed1` works without downloading another file. To try one, run `colosseum predict --weights uc3m --source path/to/image.jpg`.
65
+
66
+ Annotated images and machine-readable `predictions.json` are written to `runs/predict/`. Use `--source path/to/folder` for a folder of images, `--conf 0.25` to adjust the confidence cutoff, and `--device cuda:0` for a CUDA GPU.
67
+
68
+ To load a previously trained checkpoint in another script:
69
+
70
+ ```python
71
+ from colosseum import Colosseum
72
+
73
+ model = Colosseum("runs/custom/my_object/seed0/colosseum.pt")
74
+ for box in model.predict("path/to/image.jpg", conf=0.35):
75
+ print(box.label, box.confidence, box.x1, box.y1, box.x2, box.y2)
76
+ ```
77
+
78
+ The coordinates are pixels in the original image after EXIF orientation is applied. The prediction threshold only filters results; it does not change the model. Local confidence peaks are used to select candidate boxes.
79
+
80
+ ## Model zoo
81
+
82
+ All included weights are **final checkpoints** from the reported three-seed training runs. The top-level weights are seed 0; seed 1 and 2 are in `weights/seeds/`. They are also bundled with the pip package. The complete training state was stripped; each file contains model tensors and the metadata required for inference. The `weights/manifest.json` file records SHA-256 hashes.
83
+
84
+ | Weights | Training dataset | Detects | Output stride |
85
+ |---|---|---|---|
86
+ | [`uc3m.pt`](weights/uc3m.pt) | UC3M-LP | license plate | 2 × 8 |
87
+ | [`gtsdb.pt`](weights/gtsdb.pt) | GTSDB | prohibitory traffic sign | 4 × 4 |
88
+ | [`shwd.pt`](weights/shwd.pt) | SHWD | safety helmet | 4 × 4 |
89
+ | [`ccpd.pt`](weights/ccpd.pt) | CCPD2019 | license plate | 2 × 8 |
90
+ | [`drone.pt`](weights/drone.pt) | Drone | drone | 4 × 4 |
91
+ | [`amlpr.pt`](weights/amlpr.pt) | AMLPR | license plate | 2 × 8 |
92
+
93
+ For example, use `weights/seeds/uc3m_seed1.pt` or `weights/seeds/uc3m_seed2.pt` to select another UC3M seed. The same naming convention applies to every dataset.
94
+
95
+ The dataset name identifies the training domain; three plate checkpoints are separate models trained on different images. Choose the checkpoint closest to the images you expect. These are floating-point PyTorch weights. They are not directly loadable by ESP32-S3 firmware; device deployment requires the separately quantized INT8 program and C runtime.
96
+
97
+ ## Benchmarks
98
+
99
+ The plots below reproduce the **initial matched comparisons** in the paper: validation mean and sample SD across three seeds. The separate SHWD and Drone reruns changed two narrow initial rankings; see the paper before interpreting those datasets. The [accuracy data](assets/benchmark_results.csv), [device data](assets/benchmark_device.csv), and [plot script](scripts/plot_benchmark.py) are included.
100
+
101
+ ![Float validation AP50, AP75, and AP50:95 across six datasets](assets/benchmark_float.png)
102
+
103
+ ![Post-training INT8 validation AP50, AP75, and AP50:95 across six datasets](assets/benchmark_int8.png)
104
+
105
+ The device comparison below is for **UC3M seed 0** on a 240 MHz ESP32-S3 with an embedded frame. Training throughput is a separate RTX 5060 pilot, not a device speed measurement.
106
+
107
+ ![UC3M training throughput and ESP32-S3 deployment measurements](assets/benchmark_device.png)
108
+
109
+ The following paper figure shows selected **validation examples**, not aggregate performance. Green boxes are targets; red boxes are the top prediction when present.
110
+
111
+ ![Qualitative validation examples for Colosseum and three baselines](assets/qualitative_results.png)
112
+
113
+ ## Train on your data
114
+
115
+ Colosseum uses YOLO-format single-class annotations. For a new dataset, put images and labels in `images/train`, `images/val`, `labels/train`, and `labels/val`, then make a `data.yaml`. One command trains both stages and packages weights:
116
+
117
+ ```bash
118
+ colosseum train --data my_dataset/data.yaml --name my_object --label "my object" --grid 4x4
119
+ colosseum predict --weights runs/custom/my_object/seed0/colosseum.pt --source path/to/image.jpg
120
+ ```
121
+
122
+ The script computes box-size priors from training labels only. Use `--grid 2x8` for elongated horizontal targets, `--seeds 0 1 2` for three runs, and `--source-class N` when the target's YOLO class ID is not zero. See the [full custom-training guide](docs/TRAINING.md) for the dataset layout, options, and manual commands. Raw datasets and official test images are not included here.
123
+
124
+ ## Reproducibility
125
+
126
+ - Input: grayscale, bilinear resize preserving aspect ratio, gray padding value 114.
127
+ - Architecture: fixed-gradient stem, regular bottleneck blocks, multiscale fusion, one confidence channel and four geometry channels.
128
+ - Geometry: normalized boxes with a prior-log width and height encoding; priors are saved in each checkpoint.
129
+ - The implementation is packaged under `colosseum/`; `src/` keeps compatibility imports for older source-checkout commands.
130
+
131
+ Please cite the Colosseum paper when it becomes available. The datasets remain subject to their source terms; no raw dataset images are redistributed here.
132
+
133
+ ## License
134
+
135
+ Code and the 18 Colosseum weight files are released under the [MIT License](LICENSE). Third-party datasets are excluded.
@@ -0,0 +1,5 @@
1
+ """Compatibility imports for earlier public checkpoints and scripts."""
2
+
3
+ from .api import AKROS, Colosseum, Detection
4
+
5
+ __all__ = ["AKROS", "Colosseum", "Detection"]
@@ -0,0 +1,7 @@
1
+ """Compatibility interface for the earlier AKROS package name."""
2
+
3
+ from colosseum.api import Colosseum, Detection
4
+
5
+ AKROS = Colosseum
6
+
7
+ __all__ = ["AKROS", "Colosseum", "Detection"]
@@ -0,0 +1,5 @@
1
+ """Public Colosseum inference interface."""
2
+
3
+ from .api import Colosseum, Detection
4
+
5
+ __all__ = ["Colosseum", "Detection"]
@@ -0,0 +1,7 @@
1
+ """Allow ``python -m colosseum`` as an alternative to the console script."""
2
+
3
+ from .cli import main
4
+
5
+
6
+ if __name__ == "__main__":
7
+ main()
@@ -0,0 +1 @@
1
+ """Colosseum architecture and training components."""
@@ -0,0 +1,27 @@
1
+ """Canonical public interface for the shared MERG-Det architecture."""
2
+
3
+ from .platedet_u import (
4
+ AsymmetricResidualBlock,
5
+ ConvBNAct,
6
+ Detection,
7
+ EdgeMicroStem,
8
+ GeometryCodec,
9
+ GeometryPrior,
10
+ MergDet,
11
+ decode,
12
+ filter_by_geometry,
13
+ parameter_count,
14
+ )
15
+
16
+ __all__ = [
17
+ "AsymmetricResidualBlock",
18
+ "ConvBNAct",
19
+ "Detection",
20
+ "EdgeMicroStem",
21
+ "GeometryCodec",
22
+ "GeometryPrior",
23
+ "MergDet",
24
+ "decode",
25
+ "filter_by_geometry",
26
+ "parameter_count",
27
+ ]
@@ -0,0 +1,191 @@
1
+ """MERG-Det-M: scaled MERG-Det family (S/M/L) with a three-scale readout.
2
+
3
+ Same input path, anisotropic 2x8 output grid, five-channel head layout,
4
+ codec, loss and decoder as MERG-Det, so every existing training and
5
+ evaluation tool applies unchanged. Differences are capacity only: wider
6
+ stages, repeated blocks, a fourth downsampling stage, a top-down fusion of
7
+ three scales onto the 2x8 grid, and a decoupled two-layer head. Operators stay
8
+ conv / BN / ReLU6 / add / nearest upsample. This is an engineering variant,
9
+ not a novelty claim; the 11k MERG-Det stays a separate result.
10
+ """
11
+ from __future__ import annotations
12
+
13
+ import math
14
+ from dataclasses import dataclass
15
+
16
+ import torch
17
+ from torch import Tensor, nn
18
+ from torch.nn import functional as F
19
+
20
+ from .platedet_u import AsymmetricResidualBlock, ConvBNAct, EdgeMicroStem
21
+
22
+
23
+ @dataclass(frozen=True)
24
+ class MergDetMConfig:
25
+ name: str
26
+ channels: tuple[int, int, int, int, int] # stem, stage1, stage2, stage3, stage4
27
+ repeats: tuple[int, int, int, int] # extra non-strided blocks per stage
28
+ neck: int # lateral / fused channels on every scale
29
+ head: int # objectness branch channels
30
+ head_depth: int = 2
31
+ geometry_head: int | None = None # defaults to head for historical configs
32
+ geometry_head_depth: int | None = None # defaults to head_depth
33
+
34
+
35
+ # Budget ladder bracketing the Stage B candidates at 144x192 (single class):
36
+ # S 104,005 params / 54.4M Conv MAC (~ESPDet-Pico / 3.5)
37
+ # M 368,813 params / 129.4M Conv MAC (~ESPDet-Pico 361k, 2.7x its MAC)
38
+ # L 927,661 params / 226.0M Conv MAC (~YOLOX-Nano 897k, 2.5x its MAC)
39
+ # MACs stay above the candidates' at equal parameters because the 2x8 grid
40
+ # keeps 72 rows; this is reported, not hidden.
41
+ CONFIGS: dict[str, MergDetMConfig] = {
42
+ # V4 high-resolution narrow candidates. They retain V3's 192x256 canvas,
43
+ # repeats and 2x8 readout while reducing channel traffic enough to target
44
+ # the ESPDet/Tinyissimo latency range. V4T is the SRAM-first arm; V4S is
45
+ # the less aggressive accuracy-first arm.
46
+ "V4T": MergDetMConfig("V4T", (8, 12, 24, 56, 112), (1, 1, 2, 3), 16, 16),
47
+ "V4S": MergDetMConfig("V4S", (8, 16, 24, 48, 72), (1, 1, 2, 3), 32, 32),
48
+ # Accuracy-first rebalance after V4T's frozen pilot: spend compute at the
49
+ # input/detail path and use one wider head block; recover budget in stage3.
50
+ "V4A": MergDetMConfig("V4A", (10, 12, 24, 40, 128), (1, 1, 2, 3), 16, 24, 1),
51
+ # V4B static-screen arms: retain V4A's backbone and spend only on box
52
+ # regression, since V4A's localization diagnostic found size/vertical
53
+ # errors rather than an AP50/objectness bottleneck.
54
+ "V4B1": MergDetMConfig("V4B1", (10, 12, 24, 40, 128), (1, 1, 2, 3), 16, 16, 1, 24, 2),
55
+ "V4B2": MergDetMConfig("V4B2", (10, 12, 24, 40, 128), (1, 1, 2, 3), 16, 16, 1, 32, 2),
56
+ "V4B3": MergDetMConfig("V4B3", (10, 12, 24, 40, 128), (1, 1, 2, 3), 24, 16, 1, 32, 2),
57
+ # V4C restores V4A's objectness width after V4B1 failed its epoch-50 AP50
58
+ # gate, while retaining the geometry-only depth intervention.
59
+ "V4C1": MergDetMConfig("V4C1", (10, 12, 24, 40, 128), (1, 1, 2, 3), 16, 24, 1, 24, 2),
60
+ "V4C2": MergDetMConfig("V4C2", (10, 12, 24, 40, 128), (1, 1, 2, 3), 16, 24, 1, 32, 2),
61
+ # V5 accuracy-first/hardware-bounded arm. Restore richer multi-scale
62
+ # fusion while keeping the V4A detail path and a shallow deployable head.
63
+ "V5A": MergDetMConfig("V5A", (10, 12, 24, 40, 128), (1, 1, 2, 3), 32, 24, 1),
64
+ "S": MergDetMConfig("S", (16, 24, 32, 72, 112), (1, 1, 2, 3), 48, 48),
65
+ "M": MergDetMConfig("M", (24, 32, 48, 112, 224), (1, 2, 3, 4), 64, 64),
66
+ "L": MergDetMConfig("L", (24, 32, 64, 160, 320), (1, 2, 4, 6), 64, 64),
67
+ }
68
+
69
+ STRIDES = ((1, 2), (1, 2), (2, 2), (2, 2)) # after the stride-2 stem: grid 2x8, 4x16, 8x32
70
+
71
+ # Grid ablation (docs/GRID_ABLATION_PROTOCOL.md): the output grid is set by the
72
+ # strides of stages 1-2 only; stages 3-4 keep (2, 2) and the neck upsamples by
73
+ # size. Parameters are grid-independent by construction; MAC is not.
74
+ GRID_STRIDES: dict[tuple[int, int], tuple[tuple[int, int], tuple[int, int]]] = {
75
+ (2, 8): ((1, 2), (1, 2)), # anisotropic default: 72x24 cells at 144x192
76
+ (4, 4): ((1, 1), (2, 2)), # isotropic at equal cell count: 36x48
77
+ (8, 2): ((2, 1), (2, 1)), # wrong-direction control: 18x96
78
+ (8, 8): ((2, 2), (2, 2)), # conventional coarse grid: 18x24
79
+ }
80
+
81
+
82
+ class RegularBottleneckBlock(nn.Module):
83
+ """Hardware-oriented bottleneck using one regular 3x3 spatial convolution."""
84
+
85
+ def __init__(self, in_channels: int, out_channels: int, stride: int | tuple[int, int] = 1) -> None:
86
+ super().__init__()
87
+ stride_y, stride_x = (stride, stride) if isinstance(stride, int) else stride
88
+ self.stride_y, self.stride_x = stride_y, stride_x
89
+ hidden = max(4, min(in_channels, out_channels) // 4)
90
+ self.reduce = ConvBNAct(in_channels, hidden, 1)
91
+ self.spatial = ConvBNAct(hidden, hidden, 3, stride=(stride_y, stride_x), padding=1)
92
+ self.expand = ConvBNAct(hidden, out_channels, 1, activation=False)
93
+ if stride_y == 1 and stride_x == 1 and in_channels == out_channels:
94
+ self.skip: nn.Module = nn.Identity()
95
+ else:
96
+ self.skip = ConvBNAct(in_channels, out_channels, 1, stride=(stride_y, stride_x), activation=False)
97
+ self.activation = nn.ReLU6(inplace=True)
98
+
99
+ def forward(self, source: Tensor) -> Tensor:
100
+ residual = self.skip(source)
101
+ features = self.expand(self.spatial(self.reduce(source)))
102
+ return self.activation(features + residual)
103
+
104
+
105
+ def _stage(in_channels: int, out_channels: int, stride: tuple[int, int], repeats: int,
106
+ block_type: type[nn.Module]) -> nn.Sequential:
107
+ blocks = [block_type(in_channels, out_channels, stride=stride)]
108
+ blocks += [block_type(out_channels, out_channels) for _ in range(repeats)]
109
+ return nn.Sequential(*blocks)
110
+
111
+
112
+ class MergDetM(nn.Module):
113
+ def __init__(
114
+ self,
115
+ size: str = "M",
116
+ objectness_prior: float | None = None,
117
+ output_stride: tuple[int, int] = (2, 8),
118
+ block_kind: str = "asymmetric",
119
+ gradient_mode: str = "fixed",
120
+ fusion_mode: str = "multiscale",
121
+ ) -> None:
122
+ super().__init__()
123
+ config = CONFIGS[size]
124
+ self.config = config
125
+ output_stride = (int(output_stride[0]), int(output_stride[1]))
126
+ if output_stride not in GRID_STRIDES:
127
+ raise ValueError(f"output_stride must be one of {sorted(GRID_STRIDES)}, got {output_stride}")
128
+ self.output_stride = output_stride
129
+ if block_kind not in ("asymmetric", "regular"):
130
+ raise ValueError(f"block_kind must be 'asymmetric' or 'regular', got {block_kind!r}")
131
+ if fusion_mode not in ("multiscale", "detail_only"):
132
+ raise ValueError(f"fusion_mode must be 'multiscale' or 'detail_only', got {fusion_mode!r}")
133
+ self.block_kind = block_kind
134
+ self.gradient_mode = gradient_mode
135
+ self.fusion_mode = fusion_mode
136
+ block_type = AsymmetricResidualBlock if block_kind == "asymmetric" else RegularBottleneckBlock
137
+ stride1, stride2 = GRID_STRIDES[output_stride]
138
+ c0, c1, c2, c3, c4 = config.channels
139
+ r1, r2, r3, r4 = config.repeats
140
+ self.stem = EdgeMicroStem(c0, gradient_mode=gradient_mode)
141
+ self.stage1 = _stage(c0, c1, stride1, r1, block_type)
142
+ self.stage2 = _stage(c1, c2, stride2, r2, block_type) # output grid (detail)
143
+ self.stage3 = _stage(c2, c3, STRIDES[2], r3, block_type) # 4x16
144
+ self.stage4 = _stage(c3, c4, STRIDES[3], r4, block_type) # 8x32
145
+ neck = config.neck
146
+ self.lateral2 = ConvBNAct(c2, neck, 1, activation=False)
147
+ self.lateral3 = ConvBNAct(c3, neck, 1, activation=False)
148
+ self.lateral4 = ConvBNAct(c4, neck, 1, activation=False)
149
+ self.smooth3 = block_type(neck, neck)
150
+ self.smooth2 = block_type(neck, neck)
151
+ self.activation = nn.ReLU6(inplace=True)
152
+ geometry_head = config.geometry_head if config.geometry_head is not None else config.head
153
+ geometry_depth = (
154
+ config.geometry_head_depth
155
+ if config.geometry_head_depth is not None
156
+ else config.head_depth
157
+ )
158
+ self.objectness_branch = self._branch(neck, config.head, config.head_depth, block_type)
159
+ self.geometry_branch = self._branch(neck, geometry_head, geometry_depth, block_type)
160
+ self.objectness_head = nn.Conv2d(config.head, 1, 1)
161
+ self.geometry_head = nn.Conv2d(geometry_head, 4, 1)
162
+ if objectness_prior is not None:
163
+ if not 0.0 < objectness_prior < 1.0:
164
+ raise ValueError("objectness_prior must lie in (0, 1)")
165
+ with torch.no_grad():
166
+ self.objectness_head.bias.fill_(-math.log((1.0 - objectness_prior) / objectness_prior))
167
+
168
+ @staticmethod
169
+ def _branch(in_channels: int, channels: int, depth: int,
170
+ block_type: type[nn.Module] = AsymmetricResidualBlock) -> nn.Sequential:
171
+ # The 2x8 grid has 1728 cells at 144x192, so a dense 3x3 head would
172
+ # dominate the MAC budget; keep the head asymmetric-depthwise like the
173
+ # stages and spend parameters in the coarse stages instead.
174
+ layers = [block_type(in_channels, channels)]
175
+ layers += [block_type(channels, channels) for _ in range(depth - 1)]
176
+ return nn.Sequential(*layers)
177
+
178
+ def forward(self, image: Tensor) -> Tensor:
179
+ s2 = self.stage2(self.stage1(self.stem(image)))
180
+ if self.fusion_mode == "multiscale":
181
+ s3 = self.stage3(s2)
182
+ s4 = self.stage4(s3)
183
+ p4 = self.lateral4(s4)
184
+ p3 = self.smooth3(self.activation(self.lateral3(s3) + F.interpolate(p4, size=s3.shape[-2:], mode="nearest")))
185
+ p2 = self.smooth2(self.activation(self.lateral2(s2) + F.interpolate(p3, size=s2.shape[-2:], mode="nearest")))
186
+ else:
187
+ p2 = self.smooth2(self.activation(self.lateral2(s2)))
188
+ # objectness, dx, dy, normalized width, normalized height — MERG-Det layout
189
+ return torch.cat(
190
+ (self.objectness_head(self.objectness_branch(p2)), self.geometry_head(self.geometry_branch(p2))), dim=1
191
+ )