Single Object Detection: Training and Serving with FastAPI
Table of Contents
📚 Computer Vision Series — Session 3. Part 2 of 3. Haven't read the theory yet? Part 1 covers the concept and architecture before we get hands-on here.
Quick Recap
Part 1 laid out the big picture: an image goes through a CNN backbone, gets compressed into a handful of output numbers (is_object plus bounding box coordinates), trained with a combination of three loss functions.
Now it's time to actually build this in PyTorch with a MobileNetV2 backbone, all the way through to deploying it as an API with FastAPI. The code below is the first version that was genuinely trained and measured, not just theory on paper.
📁 Full code and the results that follow are also on GitHub if you want to run this yourself.
Setting Up the Dataset and Model
Dataset Class
The dataset reads images plus a .txt annotation file containing raw x y w h coordinates. If no annotation file exists, that image is treated as having no object.
class ObjectDetectionDataset(Dataset):
def __init__(self, image_folder, transform=None):
self.image_folder = image_folder
self.image_files = [f for f in os.listdir(image_folder) if f.endswith(('png', 'jpg', 'jpeg'))]
self.transform = transform
def __len__(self):
return len(self.image_files)
def __getitem__(self, idx):
img_name = self.image_files[idx]
img_path = os.path.join(self.image_folder, img_name)
annotation_path = os.path.join(self.image_folder, img_name.replace('jpg', 'txt'))
image = Image.open(img_path).convert("RGB")
img_width, img_height = image.size
is_object = torch.tensor([0], dtype=torch.float32)
bbox = torch.tensor([0, 0, 0, 0], dtype=torch.float32)
if os.path.exists(annotation_path):
with open(annotation_path, 'r') as f:
lines = f.readlines()
if len(lines) > 0:
x, y, w, h = map(float, lines[0].strip().split())
x /= img_width; y /= img_height; w /= img_width; h /= img_height
is_object = torch.tensor([1], dtype=torch.float32)
bbox = torch.tensor([x, y, w, h], dtype=torch.float32)
if self.transform:
image = self.transform(image)
return image, is_object, bbox
One detail worth flagging: bounding box coordinates are normalized to a 0-1 range, not raw pixels. That keeps the loss scale consistent regardless of image size.
Model: MobileNetV2 with a Custom Head
class ObjectDetectionModel(nn.Module):
def __init__(self):
super(ObjectDetectionModel, self).__init__()
self.backbone = mobilenet_v2(weights="DEFAULT").features
self.classifier = nn.Sequential(
nn.AdaptiveAvgPool2d((1, 1)),
nn.Flatten(),
nn.Linear(1280, 512),
nn.ReLU(),
nn.Linear(512, 5) # 1 is_object + 4 bbox
)
def forward(self, x):
x = self.backbone(x)
x = self.classifier(x)
is_object = torch.sigmoid(x[:, :1])
bbox = torch.sigmoid(x[:, 1:])
return is_object, bbox
Small aside: the original source material used mobilenet_v2(pretrained=True), which is deprecated in newer torchvision releases. Swapping it for weights="DEFAULT" keeps it working with current environments. Little things like this tend to surface the moment old code gets run again.
AdaptiveAvgPool2d((1, 1)) here flattens the entire feature map into a single vector before the fully connected layer. Simple, and good enough for this first version, but there's a trade-off buried in that choice that only becomes visible later. More on that in Part 3.
IoU: Measuring How Good a Prediction Is
IoU (Intersection over Union) measures how well a predicted bounding box overlaps with the ground truth. It's not used for training, that's the loss function's job, but for evaluating how good the results actually are.
def calculate_iou(box1, box2):
x1, y1, w1, h1 = box1
x2, y2, w2, h2 = box2
x_left = max(x1 - w1/2, x2 - w2/2)
y_top = max(y1 - h1/2, y2 - h2/2)
x_right = min(x1 + w1/2, x2 + w2/2)
y_bottom = min(y1 + h1/2, y2 + h2/2)
if x_right < x_left or y_bottom < y_top:
return 0.0
intersection = (x_right - x_left) * (y_bottom - y_top)
union = w1*h1 + w2*h2 - intersection
return intersection / union
The combined loss (MixedLoss) merges Binary Cross Entropy for objectness with MSE for location, with tunable weights:
class MixedLoss(nn.Module):
def __init__(self, weight_objectness=1.0, weight_localization=1.0):
super(MixedLoss, self).__init__()
self.bce_loss = nn.BCELoss()
self.mse_loss = nn.MSELoss()
self.weight_o = weight_objectness
self.weight_l = weight_localization
def forward(self, pred_object, true_object, pred_bbox, true_bbox):
loss_object = self.bce_loss(pred_object, true_object)
loss_bbox = self.mse_loss(pred_bbox, true_bbox)
return self.weight_o * loss_object + self.weight_l * loss_bbox
🤔 Quick check: if IoU is only used for evaluation and not training, why not just use it as the loss function directly? Click for the answer.
Good instinct, and it turns into one of the main topics in Part 3. This first version doesn't use it as a loss yet, only as a metric tracked outside the training loop.
The Full Training Loop
def train_model(model, train_loader, test_loader, criterion, optimizer, device, num_epochs=10):
model.train()
best_iou = 0.0
for epoch in range(num_epochs):
running_loss, total_iou_train, total_iou_test = 0.0, 0.0, 0.0
for images, is_object, bboxes in train_loader:
images, is_object, bboxes = images.to(device), is_object.to(device), bboxes.to(device)
optimizer.zero_grad()
pred_object, pred_bbox = model(images)
loss = criterion(pred_object, is_object, pred_bbox, bboxes)
loss.backward()
optimizer.step()
running_loss += loss.item()
for i in range(len(bboxes)):
if is_object[i] > 0.5:
total_iou_train += calculate_iou(pred_bbox[i].detach().cpu().numpy(), bboxes[i].cpu().numpy())
avg_iou_train = total_iou_train / len(train_loader.dataset)
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}, IoU (Train): {avg_iou_train:.4f}")
model.eval()
with torch.no_grad():
for images, is_object, bboxes in test_loader:
images, is_object, bboxes = images.to(device), is_object.to(device), bboxes.to(device)
pred_object, pred_bbox = model(images)
for i in range(len(bboxes)):
if is_object[i] > 0.5:
total_iou_test += calculate_iou(pred_bbox[i].cpu().numpy(), bboxes[i].cpu().numpy())
avg_iou_test = total_iou_test / len(test_loader.dataset)
print(f"IoU (Test): {avg_iou_test:.4f}")
torch.save(model.state_dict(), 'last.pt')
if avg_iou_test > best_iou:
best_iou = avg_iou_test
torch.save(model.state_dict(), 'best.pt')
model.train()
First Training Results
The dataset has 241 car images, split 192 for training and 49 for testing (80/20). Training ran on Google Colab with a Tesla T4 GPU for 30 epochs.
| Metric | Value |
|---|---|
| Best IoU (test) | 0.7186, at epoch 27 |
| IoU (train, final epoch) | 0.8382 |
| Train/test gap | 0.120 |
These numbers matter because they're real, not estimated. A 0.120 gap between train and test IoU points to mild overfitting: the model is a bit more "familiar" with the training data than its actual ability to generalize. That gap turned out to be a useful starting point, and it's exactly what Part 3 digs into.
Serving the Model with FastAPI
The trained model gets wrapped into a simple API:
import torch
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from torchvision import transforms
from torchvision.models import mobilenet_v2
from torch import nn
from PIL import Image
import uvicorn
import io
app = FastAPI()
model = ObjectDetectionModel()
model.load_state_dict(torch.load("best.pt", map_location=torch.device('cpu')))
model.eval()
transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
try:
image = Image.open(io.BytesIO(await file.read())).convert("RGB")
image_tensor = transform(image).unsqueeze(0)
with torch.no_grad():
is_object, bbox = model(image_tensor)
is_object_value = is_object.item()
bbox = bbox.squeeze(0).tolist()
if is_object_value > 0.5:
response = {"is_object": True, "bbox": {"x": bbox[0], "y": bbox[1], "w": bbox[2], "h": bbox[3]}}
else:
response = {"is_object": False, "bbox": None}
return JSONResponse(content=response)
except Exception as e:
return JSONResponse(status_code=500, content={"error": str(e)})
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)
The /predict endpoint takes an image via multipart upload and returns JSON with is_object and bbox. If you're deploying straight from Colab without your own server, there's also a variant using pyngrok for tunneling plus bounding box visualization drawn directly on the image. Full code is in the repo.
🧠 Quiz time: click to open
1. Why does the API return bbox: None when is_object is false, instead of coordinates like {0,0,0,0}? Because {0,0,0,0} could be misread as a valid bounding box sitting at the origin. None explicitly says there's no box to show, which is safer for whatever client consumes this API.
2. Why does this endpoint use async def for the predict function? Reading the uploaded file (await file.read()) is an I/O operation that ideally shouldn't block other incoming requests while it waits. FastAPI supports this natively through async.
Cheat Sheet
[x] Dataset reads images plus
.txtannotations, normalizes bbox to 0-1[x] Model: MobileNetV2 backbone (transfer learning) plus a custom head
[x] Loss: BCE (objectness) + MSE (localization), weighted
[x] IoU used as an evaluation metric, not a loss
[x] v1 results: test IoU 0.7186, train/test gap 0.120
[x] Serving: FastAPI
/predictendpoint, takes an image, returns JSON
Coming Up in Part 3
The v1 model above "works," but testing it on real-world photos outside the training set surfaced a pattern nobody expected. Part 3 covers two more rounds of iteration, and one finding that changes how you should read evaluation numbers from a small model like this one.
📁 Dataset: Google Drive 💻 Repo: GitHub
Part of the Computer Vision Super Class series, Session 3: Single Object Detection.
