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.
- colosseum_microcontroller-0.1.3/LICENSE +21 -0
- colosseum_microcontroller-0.1.3/PKG-INFO +54 -0
- colosseum_microcontroller-0.1.3/PYPI_README.md +36 -0
- colosseum_microcontroller-0.1.3/README.md +135 -0
- colosseum_microcontroller-0.1.3/akros/__init__.py +5 -0
- colosseum_microcontroller-0.1.3/akros/api.py +7 -0
- colosseum_microcontroller-0.1.3/colosseum/__init__.py +5 -0
- colosseum_microcontroller-0.1.3/colosseum/__main__.py +7 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/__init__.py +1 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/mergdet.py +27 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/mergdet_m.py +191 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/plate_dataset.py +354 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/plate_loss.py +647 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/plate_metrics.py +182 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/platedet_u.py +429 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/train_platedet.py +538 -0
- colosseum_microcontroller-0.1.3/colosseum/_core/training_state.py +96 -0
- colosseum_microcontroller-0.1.3/colosseum/api.py +159 -0
- colosseum_microcontroller-0.1.3/colosseum/cli.py +22 -0
- colosseum_microcontroller-0.1.3/colosseum/evaluation.py +68 -0
- colosseum_microcontroller-0.1.3/colosseum/pack_checkpoint.py +50 -0
- colosseum_microcontroller-0.1.3/colosseum/prediction.py +54 -0
- colosseum_microcontroller-0.1.3/colosseum/training.py +170 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/amlpr.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/ccpd.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/drone.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/gtsdb.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/manifest.json +152 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/amlpr_seed1.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/amlpr_seed2.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/ccpd_seed1.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/ccpd_seed2.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/drone_seed1.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/drone_seed2.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/gtsdb_seed1.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/gtsdb_seed2.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/shwd_seed1.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/shwd_seed2.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/uc3m_seed1.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/seeds/uc3m_seed2.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/shwd.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum/weights/uc3m.pt +0 -0
- colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/PKG-INFO +54 -0
- colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/SOURCES.txt +48 -0
- colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/dependency_links.txt +1 -0
- colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/entry_points.txt +2 -0
- colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/requires.txt +3 -0
- colosseum_microcontroller-0.1.3/colosseum_microcontroller.egg-info/top_level.txt +2 -0
- colosseum_microcontroller-0.1.3/pyproject.toml +33 -0
- 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&logoColor=white"></a>
|
|
11
|
+
<a href="#model-zoo"><img alt="PyTorch models" src="https://img.shields.io/badge/PyTorch-models-EE4C2C?logo=pytorch&logoColor=white"></a>
|
|
12
|
+
<a href="#benchmarks"><img alt="ESP32-S3 benchmark" src="https://img.shields.io/badge/ESP32--S3-benchmark-00979D?logo=espressif&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> ·
|
|
18
|
+
<a href="notebooks/Train_Colosseum.ipynb">📓 Training notebook</a> ·
|
|
19
|
+
<a href="#model-zoo">🧠 Models</a> ·
|
|
20
|
+
<a href="#benchmarks">📊 Benchmarks</a> ·
|
|
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
|
+

|
|
102
|
+
|
|
103
|
+

|
|
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
|
+

|
|
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
|
+

|
|
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 @@
|
|
|
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
|
+
)
|