Writing your first AI model with PyTorch and Torchvision
I had video from the range and wanted a model that could find the ball in each frame. This is the loop I ended up with: a Faster R-CNN from Torchvision doing the guessing, and me telling it when it was wrong.
Two files carry the whole thing. ball_analysis.py runs the loop: it picks random frames from the video with OpenCV, runs the detector, and asks me for feedback. ball_calibration.json is where every correction lands.
The loop
For each frame the model proposes a ball position, and I do one of three things: confirm it, mark that there is no ball in the frame, or click where the ball actually is. Each answer is appended to the calibration file. Sampling frames at random mattered more than I expected. Going through the video sequentially teaches the model a lot about one swing and nothing about the rest of the session.
When enough frames have been calibrated the loop breaks and everything is saved for retraining.
What the calibration data looks like
Every entry records the frame number, where the ball actually was (or that none was visible), the model’s confidence, and a timestamp:
[
{
"frame_number": 173401,
"x": 1560.0,
"y": 372.0,
"timestamp": 1737929567.2006788,
"ai_confidence": 0.7186214327812195,
"ai_validated": true,
"is_correction": false,
"no_ball": false,
"session_id": "20250126_231247",
"ai_original_prediction": null
},
{
"frame_number": 295534,
"x": null,
"y": null,
"timestamp": 1737929570.059416,
"ai_confidence": null,
"ai_validated": false,
"is_correction": true,
"no_ball": true,
"session_id": "20250126_231250",
"ai_original_prediction": {
"x": null,
"y": null,
"confidence": 0.0
}
},
{
"frame_number": 9496,
"x": null,
"y": null,
"timestamp": 1737929573.103283,
"ai_confidence": null,
"ai_validated": false,
"is_correction": true,
"no_ball": true,
"session_id": "20250126_231253",
"ai_original_prediction": {
"x": null,
"y": null,
"confidence": 0.0
}
},
{
"frame_number": 18141,
"x": 1689.0,
"y": 388.0,
"timestamp": 1737929575.656329,
"ai_confidence": 0.41590890288352966,
"ai_validated": true,
"is_correction": false,
"no_ball": false,
"session_id": "20250126_231255",
"ai_original_prediction": null
},
{
"frame_number": 264335,
"x": null,
"y": null,
"timestamp": 1737929578.010589,
"ai_confidence": null,
"ai_validated": false,
"is_correction": true,
"no_ball": true,
"session_id": "20250126_231258",
"ai_original_prediction": {
"x": null,
"y": null,
"confidence": 0.0
}
}
]
Together these entries become the ground truth the model learns from. The no_ball corrections turned out to be as valuable as the position clicks. The early model saw balls everywhere.
Why the human stays in the loop
Quick retraining steps use the corrections you just made, so the model improves during a session. Offline, the full calibration history makes for a stronger retrain. You steer the model by telling it when it is right and wrong. There is no labelling marathon up front, just corrections where they are needed.
Two files talking to each other is enough to start. The model gets more reliable every session, and the same JSON keeps working as the dataset grows.