# Single Object Detection: Training and Serving with FastAPI

## Table of Contents

*   [Quick Recap](#quick-recap)
    
*   [Setting Up the Dataset and Model](#setting-up-the-dataset-and-model)
    
*   [IoU: Measuring How Good a Prediction Is](#iou-measuring-how-good-a-prediction-is)
    
*   [The Full Training Loop](#the-full-training-loop)
    
*   [First Training Results](#first-training-results)
    
*   [Serving the Model with FastAPI](#serving-the-model-with-fastapi)
    
*   [Cheat Sheet](#cheat-sheet)
    
*   [Coming Up in Part 3](#coming-up-in-part-3)
    

> 📚 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](https://github.com/arielshakaramiro/single-object-detection-car) 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.

```python
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

```python
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.

```python
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:

```python
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

```python
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:

```python
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 `.txt` annotations, 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 `/predict` endpoint, 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](https://drive.google.com/drive/folders/13qoW2HCHWUVfzr4EFKUc3ZngoQT_D97A?usp=sharing) 💻 Repo: [GitHub](https://github.com/arielshakaramiro/single-object-detection-car)

* * *

*Part of the Computer Vision Super Class series, Session 3: Single Object Detection.*