https://github.com/anto18671/coco-pretraining
Train a Mask R-CNN instance segmentation model on the COCO 2017 dataset using a timm ConvNeXt-Tiny backbone.
https://github.com/anto18671/coco-pretraining
coco machine-learning mask-rcnn pretraining segmentation
Last synced: 9 days ago
JSON representation
Train a Mask R-CNN instance segmentation model on the COCO 2017 dataset using a timm ConvNeXt-Tiny backbone.
- Host: GitHub
- URL: https://github.com/anto18671/coco-pretraining
- Owner: anto18671
- License: mit
- Created: 2026-05-13T13:11:14.000Z (2 months ago)
- Default Branch: main
- Last Pushed: 2026-05-13T13:50:05.000Z (2 months ago)
- Last Synced: 2026-06-03T04:33:42.753Z (about 2 months ago)
- Topics: coco, machine-learning, mask-rcnn, pretraining, segmentation
- Language: Python
- Homepage:
- Size: 6.84 KB
- Stars: 1
- Watchers: 0
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# COCO Mask R-CNN Pretraining
Train a Mask R-CNN instance segmentation model on the COCO 2017 dataset using a `timm` ConvNeXt-Tiny backbone.
This project loads COCO image annotations with `pycocotools`, builds a TorchVision `MaskRCNN` model, trains on the COCO train split, validates on the COCO val split, and saves latest, epoch, best, and final checkpoints.
## Features
- COCO 2017 instance segmentation training
- ConvNeXt-Tiny backbone from `timm`
- TorchVision Mask R-CNN detection head
- Train and validation loss tracking
- Progress bars with `tqdm`
- Automatic checkpoint saving
- Basic dataset path validation with clear file errors
## Project Structure
```text
coco-pretraining/
train.py
README.md
LICENSE
.gitignore
data/
annotations_trainval2017/
annotations/
instances_train2017.json
instances_val2017.json
train2017/
train2017/
*.jpg
val2017/
val2017/
*.jpg
```
## Dataset
Dataset used: [COCO 2017 from Kaggle](https://www.kaggle.com/datasets/sabahesaraki/2017-2017)
After downloading and extracting the dataset, make sure the folders match the layout above. The training script expects these paths:
```text
data/annotations_trainval2017/annotations/instances_train2017.json
data/annotations_trainval2017/annotations/instances_val2017.json
data/train2017/train2017/
data/val2017/val2017/
```
The `data/` folder is ignored by Git because the dataset is large.
## Installation
Create and activate a virtual environment:
```bash
python -m venv .venv
```
On Windows PowerShell:
```powershell
.\.venv\Scripts\Activate.ps1
```
Install the required packages:
```bash
pip install torch torchvision timm opencv-python numpy pycocotools tqdm
```
For GPU training, install the PyTorch build that matches your CUDA version from the official PyTorch installation guide.
## Training
Run:
```bash
python train.py
```
Current training configuration in `train.py`:
```text
Backbone: convnext_tiny
Image size: 512 x 512
Batch size: 6
Epochs: 10
Learning rate: 1e-4
Workers: 3
Classes: 91
Device: CUDA if available, otherwise CPU
```
You can adjust these values near the top of `train.py`.
## Outputs
The script saves checkpoints in the project root:
```text
maskrcnn_latest.pth
maskrcnn_epoch_1.pth
maskrcnn_epoch_2.pth
...
maskrcnn_best.pth
maskrcnn_final.pth
```
`maskrcnn_best.pth` is saved when validation loss improves. `maskrcnn_final.pth` contains the final model state dictionary after training completes.
## Notes
- Full COCO training is compute intensive and is strongly recommended on a CUDA-capable GPU.
- COCO category IDs go up to 90, so the script uses `NUM_CLASSES = 91`.
- Images are resized to `512 x 512` before training.
- Checkpoints can become large, so avoid committing `.pth` files to GitHub.
## License
This project is licensed under the terms of the included [LICENSE](LICENSE).