A console/terminal pipeline for object detection with ResNet50-SSD512 (Single Shot MultiBox Detector, 512×512 input, PyTorch), supporting Windows 11 and Ubuntu Linux 24.04. The engine (jinfagang/ssds_pytorch) lives under src/ssds/ and is driven by a per-dataset YAML config (data/{dataset}/ssd_resnet50_train_voc.yml). The example dataset steel contains steel-surface defect images in 4 classes (inclusion, patches, rolled_in_scale, scratches) and ships with a checkpoint trained to epoch 20000, so inference works out of the box, on the CPU too. Training resumes from that checkpoint and realistically wants an NVIDIA GPU; the progress is viewable in TensorBoard.
[Folder Structure]
- CLAUDE.md / AGENT.md: program specification (identical content).
- CHANGELOG.md: change log.
- copyright.txt: copyright notice.
- 1512.02325.pdf: SSD reference paper.
- data/{dataset}/labels.txt: __background__ followed by one class name per line, maintained by hand.
- data/{dataset}/ssd_resnet50_train_voc.yml: the per-dataset training and inference config (network, epochs, batch size, learning rate and the paths).
- data/{dataset}/train/ and val/: each holds images/, annotations/ (Pascal VOC XML) and the train.txt or val.txt written by step 5 (one file name per line).
- data/{dataset}/test/images/: the test images used for inference.
- model/{dataset}/: ssd_resnet_50_voc_epoch_{N}.pth (the checkpoints), checkpoint_list.txt (the training history, used to resume), logs/ (TensorBoard) and output/ (the annotated inference results).
- src/ssds/: the engine (train.py, demo.py, inference_folder.py, lib/).
- src/annotation/: the bundled aiohttp annotation server used by step 4.
- tool/tools.py: shared helpers.
- src/: Python entry points, launcher scripts and this document.
- venv/: Python virtual environment, created by step 01.
[Prerequisites]
- Operating system: Windows 11 or Ubuntu Linux 24.04.
- Python 3 and pip.
- An NVIDIA GPU is recommended for training; inference also runs on CPU.
[Instructions]
Step 1: Create the virtual environment (CPU)
Delete the old venv/, create a fresh one and install the packages listed in src/requirements.txt. The script aborts loudly if pip fails.
Windows: src\01_venv_create.bat
Linux: bash src/01_venv_create.sh
Step 2: Create the virtual environment (GPU)
Same as step 1, but swaps torch for the CUDA 12.8 (cu128) build, for machines with an NVIDIA GPU (driver CUDA version 12.8 or later, check with nvidia-smi).
Windows: src\01_venv_create_gpu.bat
Linux: bash src/01_venv_create_gpu.sh
Step 3: Activate the virtual environment
Activate the venv. On Linux, source the script.
Windows: src\02_venv_activate.bat
Linux: source src/02_venv_activate.sh
Step 4: Annotate the images
Launch the annotation web UI on port 8801 to label the train, val and test images. The annotations are written as Pascal VOC XML into annotations/.
Windows: src\03_annotation.bat
Linux: bash src/03_annotation.sh
Step 5: Write the file lists
Scan the image folders and write train.txt and val.txt (one image file name per line).
Windows: src\04_prepare.bat
Linux: bash src/04_prepare.sh
Step 6: Train the model
Train according to the dataset's YAML config. Training resumes from the last checkpoint recorded in checkpoint_list.txt and detects the GPU automatically. Press Ctrl+C at any time to stop.
Windows: src\05_train.bat
Linux: bash src/05_train.sh
Step 7: Watch the training curves
Open the logs under model/{dataset}/logs/ with TensorBoard on port 6066. Press Ctrl+C to stop it.
Windows: src\06_tensorboard.bat
Linux: bash src/06_tensorboard.sh
Step 8: Detect in a single image
Detect objects in one image, by default data/steel/test/images/patches-94.jpg.
Windows: src\07_inference.bat
Linux: bash src/07_inference.sh
Step 9: Detect in a whole folder
Detect objects in every image of a folder, the test split by default.
Windows: src\08_inference_folder.bat
Linux: bash src/08_inference_folder.sh
Step 10: Remove the virtual environment
Delete the venv/ folder.
Windows: src\99_venv_clear.bat
Linux: bash src/99_venv_clear.sh
[Python Entry Points]
- src/annotation_server.py: --dataset steel --split train|val|test --port 8801 [--no-browser]. The bundled Pascal-VOC-XML annotation UI; it writes the XML into data/{dataset}/{split}/annotations/.
- src/prepare.py: --dataset steel. Writes train.txt and val.txt (one image file name per line). labels.txt is maintained by hand.
- src/train.py: --dataset steel [--cfg YML] [--clear-logs]. Everything else comes from the YAML config (as shipped: batch 1, SGD learning rate 0.00005, SGDR scheduler, MAX_EPOCHS 30000, a checkpoint every 100 epochs, and both the train and eval phases). Training resumes from the last entry in model/{dataset}/checkpoint_list.txt (as shipped, epoch 20000). Stop with Ctrl+C whenever the loss is good enough.
- src/tensorboard_server.py: --dataset steel --port 6066 [--no-browser]. TensorBoard for the training curves.
- src/inference.py: --dataset steel [--source IMAGE] [--cfg YML] [--show]. The annotated image lands in model/{dataset}/output/{name}_result.jpg. Runs on the CPU when no GPU is present.
- src/inference_folder.py: --dataset steel [--source FOLDER] [--cfg YML] [--show]. Detects in every .jpg and .png of the folder; the results land in model/{dataset}/output/.
[Using Your Own Dataset]
The full workflow is: label the train and val images with step 4, edit data/{dataset}/labels.txt and copy and adjust the YAML config (NUM_CLASSES is the class count plus the background, along with the paths), then run step 5 to write the file lists, step 6 to train (watch it with the TensorBoard from step 7 and stop with Ctrl+C when the loss is good enough), and finally steps 8 and 9 for inference.
[GPU]
The default PyPI torch wheel is CPU-only on Windows. For an NVIDIA GPU, install the build matching your driver's CUDA version (check with nvidia-smi), for example CUDA 12.8:
pip uninstall -y torch torchvision
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
Or simply run 01_venv_create_gpu from step 2. The engine auto-detects CUDA and falls back to the CPU.
[References]
- SSD paper (Liu et al. 2016, shipped as 1512.02325.pdf): https://arxiv.org/abs/1512.02325
- ssds_pytorch (engine base): https://github.com/jinfagang/ssds_pytorch
- PyTorch install: https://pytorch.org/get-started/locally/
This SDK is built in AppForAI - AI Dev Tools.
Purchase license separately: USD 600, permanent authorization, single APP authorization, single machine authorization, one-year activation, one-year download, one-year update, one-year email technical support.