refactor api

api/.dockerignore

@@ -1,9 +1,13 @@
-# Allowlist: the Dockerfile only copies these three things, so nothing
+# Allowlist: only what the server needs is sent to the Docker daemon, so
-# else (training code, datasets, saved models, caches) is sent to the daemon.
+# training code, datasets, saved models and caches stay out of the image.
 *
 !requirements.txt
-!src/production
+!start_server.sh
 !build_models
+!src/__init__.py
+!src/common
+!src/models
+!src/production
 
 # Re-exclude junk inside the allowed paths
 **/__pycache__

api/Dockerfile

@@ -12,8 +12,11 @@ COPY requirements.txt .
 
 RUN pip install --no-cache-dir -r requirements.txt
 
-COPY src/production/. .
+# Only the code the server needs (see .dockerignore): src/production plus the
+# shared src/models and src/common packages.
+COPY src ./src
 COPY build_models ./build_models
+COPY start_server.sh .
 
 RUN chmod +x start_server.sh
 

api/src/bird_cnn/bird_cnn.py

@@ -1,152 +0,0 @@
-import os
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from tqdm import tqdm
-
-class  SeparableConvolution(nn.Module):
-    def __init__(self, c_in, c_out, kernel_size):
-        super().__init__()
-        self.depthwise = nn.Conv2d(c_in, c_in, kernel_size, groups=c_in, padding=kernel_size//2)
-        self.bn1 = nn.BatchNorm2d(c_in)
-        self.pointwise = nn.Conv2d(c_in, c_out, kernel_size=1)
-        self.bn2 = nn.BatchNorm2d(c_out)
-
-    def forward(self, x):
-        x = self.depthwise(x)
-        x = self.bn1(x)
-        x = F.relu(x)
-
-        x = self.pointwise(x)
-        x = self.bn2(x)
-        x = F.relu(x)
-
-        return x
-    
-class SkipBlock(nn.Module):
-    def __init__(self, c_in, c_out, kernel_size=3):
-        super().__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size//2),
-            nn.BatchNorm2d(c_out),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
-            nn.BatchNorm2d(c_out),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
-            nn.BatchNorm2d(c_out),
-            nn.ReLU(inplace=True)
-        )
-        self.conv_skip = nn.Sequential(
-            nn.Conv2d(c_in, c_out, 1),
-            nn.BatchNorm2d(c_out),
-            nn.ReLU(inplace=True)
-        )
-
-    def forward(self, x):
-        return(F.relu(self.conv_skip(x) + self.conv(x), inplace=True))
-
-
-class Bird_CNN(nn.Module):
-    def __init__(self, c_in, c_hidden, c_out):
-        super().__init__()
-
-        self.model = nn.Sequential(
-            nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
-            nn.BatchNorm2d(c_hidden),
-            nn.ReLU(inplace=True),
-
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            
-            SkipBlock(c_in=c_hidden, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-
-            nn.Conv2d(c_hidden*2, c_hidden*4, kernel_size=3, padding=1),
-            nn.ReLU(inplace=True),
-
-            nn.AdaptiveAvgPool2d((1, 1)),
-            nn.Flatten(),
-            nn.Linear(c_hidden*4, c_out),
-            nn.Dropout(0.3)
-        )
-
-    def forward(self, x):
-        return self.model(x)
-
-
-def trainCNN(model, optimizer, loss_module, train_data_loader, validation_data_loader, device, num_epochs, SAVE_PATH, save=False):
-
-    best_val = 0
-    for epoch in range(num_epochs):
-        ############
-        # Training #
-        ############
-        model.train()
-
-        true_preds, count = 0, 0
-        for data_inputs, classes in tqdm(train_data_loader, desc=f"Train Epoch {epoch+1}", leave=False):
-            data_inputs = data_inputs.to(device)
-            classes = classes.to(device)
-
-            preds = model(data_inputs)
-
-            loss = loss_module(preds, classes)
-
-            optimizer.zero_grad()
-
-            loss.backward()
-
-            optimizer.step()
-
-            true_preds += (preds.argmax(dim=1) == classes).sum().item()
-            count += data_inputs.size(0)
-        train_acc = true_preds / count
-
-        torch.cuda.empty_cache()
-
-        ##############
-        # Validation #
-        ##############
-        model.eval()
-
-        true_preds, count = 0, 0
-        for data_inputs, classes in tqdm(validation_data_loader, desc=f"Validate Epoch {epoch+1}", leave=False):
-            with torch.no_grad():
-                data_inputs = data_inputs.to(device)
-                classes = classes.to(device)
-
-                preds = model(data_inputs)
-
-                loss = loss_module(preds, classes)
-
-                true_preds += (preds.argmax(dim=1) == classes).sum().item()
-                count += data_inputs.size(0)
-        val_acc = true_preds / count
-
-        if(save and best_val < val_acc):
-            best_val = val_acc
-            save_dir = os.path.join(SAVE_PATH, "bird_cnn")
-            os.makedirs(save_dir, exist_ok=True)
-
-            save_path = os.path.join(save_dir, f"bird_cnn")
-            torch.save(model.state_dict(), save_path)
-
-        print(f"epoch: {epoch+1} | train accuracy: {int(train_acc * 1000) / 10}% | validation accuracy: {int(val_acc * 1000) / 10}%")
-        torch.cuda.empty_cache()
-
-
-def sample(model, img, device, SAVE_PATH, model_name="bird_cnn", folder="bird_cnn"):
-    with torch.no_grad():
-        full_path = os.path.join(SAVE_PATH, folder, model_name)
-        state_dict = torch.load(full_path, weights_only=False)
-        model.load_state_dict(state_dict)
-        model.eval()
-        img = img.to(device)
-        pred = model(img)
-        probs = F.softmax(pred, dim=1)
-        return(torch.max(probs, dim=1))

api/src/bird_cnn/config.py

@@ -0,0 +1,33 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Config:
+    """Settings for the bird classifier. Paths are relative to api/.
+    Every field can be overridden on the command line, e.g. --epochs 20."""
+
+    # Data
+    data_dir: str = "data/CUB_200_2011/images"   # ImageFolder layout: one folder per class
+    train_fraction: float = 0.8
+    image_size: int = 128
+
+    # Model
+    num_classes: int = 200
+    c_hidden: int = 16
+
+    # Training
+    epochs: int = 500
+    batch_size: int = 64
+    learning_rate: float = 1e-3
+    weight_decay: float = 1e-4
+    num_workers: int = 0
+    device: str = "auto"                          # "auto", "cpu", "cuda:0", ...
+
+    # Output
+    save_dir: str = "saved_models"
+    model_name: str = "bird_cnn"
+    save: bool = True
+
+    # predict task
+    image: str = "test1.jpg"
+    checkpoint: str = ""                          # empty = <save_dir>/<model_name>/<model_name>

api/src/bird_cnn/tasks.py

@@ -0,0 +1,87 @@
+"""Bird classifier tasks. Run with: python -m src.main bird.<task>"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+
+from src.bird_cnn.config import Config
+from src.common.checkpoints import checkpoint_path, load_checkpoint
+from src.common.data import split_dataset
+from src.common.device import get_device
+from src.common.training import fit
+from src.common.visualize import show_batch
+from src.models.bird_cnn import BirdSpecies, Bird_CNN
+
+
+def eval_transform(cfg: Config):
+    return transforms.Compose([
+        transforms.Resize(cfg.image_size),
+        transforms.CenterCrop(cfg.image_size),
+        transforms.ToTensor()
+    ])
+
+
+def train_transform(cfg: Config):
+    return transforms.Compose([
+        transforms.RandomAffine(degrees=35, translate=(0.2, 0.2)),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.Resize(cfg.image_size),
+        transforms.CenterCrop(cfg.image_size),
+        transforms.ToTensor()
+    ])
+
+
+def build_loaders(cfg: Config):
+    dataset = datasets.ImageFolder(cfg.data_dir)
+    train_ds, val_ds = split_dataset(dataset, cfg.train_fraction, train_transform(cfg), eval_transform(cfg))
+    train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers)
+    val_loader = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers)
+    return train_loader, val_loader
+
+
+def build_model(cfg: Config):
+    return Bird_CNN(c_in=3, c_hidden=cfg.c_hidden, c_out=cfg.num_classes)
+
+
+def train(cfg: Config):
+    """Train the classifier; saves the model with the best validation accuracy."""
+    device = get_device(cfg.device)
+    train_loader, val_loader = build_loaders(cfg)
+
+    model = build_model(cfg).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
+
+    fit(model, nn.CrossEntropyLoss(), optimizer, train_loader, val_loader, device,
+        epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.model_name, save=cfg.save,
+        monitor="acc", track_accuracy=True)
+
+
+def view_data(cfg: Config):
+    """Show a few augmented training images with their labels."""
+    train_loader, _ = build_loaders(cfg)
+    show_batch(train_loader)
+
+
+def predict(cfg: Config):
+    """Classify one image (--image) with a saved model."""
+    device = get_device(cfg.device)
+    path = cfg.checkpoint or checkpoint_path(cfg.save_dir, cfg.model_name)
+    model = load_checkpoint(build_model(cfg), path, device)
+
+    image = eval_transform(cfg)(Image.open(cfg.image).convert("RGB")).unsqueeze(0).to(device)
+    with torch.inference_mode():
+        confidence, pred = torch.max(F.softmax(model(image), dim=1), dim=1)
+
+    cls = pred.item()
+    name = BirdSpecies(cls).name if cfg.num_classes == len(BirdSpecies) else f"class {cls}"
+    print(f"Species: {name} | Confidence: {confidence.item():.2f}")
+
+
+TASKS = {
+    "train": train,
+    "view-data": view_data,
+    "predict": predict,
+}

api/src/bird_cnn/train_bird.py

@@ -1,72 +0,0 @@
-from enum import Enum
-
-import torch
-import torch.nn as nn
-import torchvision
-from PIL import Image
-from torchvision import datasets, transforms
-from torch.utils.data import DataLoader, random_split
-
-from src.bird_cnn.bird_cnn import Bird_CNN, sample, trainCNN
-from src.util import TransformedSubset, visualizeData
-
-SAVE_PATH = "../saved_models"
-IMAGE_SIZE = 128
-
-class bird_species(Enum):
-    Common_Kingfisher = 0
-    Common_Myna = 1
-    House_Crow = 2
-    Indian_Peacock = 3
-    Indian_Pitta = 4
-    Ruddy_Shelduck = 5
-    Sarus_Crane = 6
-
-transform_augemnt = transforms.Compose([
-    transforms.RandomAffine(
-        degrees=35,                # no rotation
-        translate=(0.2, 0.2)      # shift up to 20% horizontally/vertically
-    ),
-    transforms.RandomHorizontalFlip(p=0.5),
-    transforms.Resize(IMAGE_SIZE),
-    transforms.CenterCrop(IMAGE_SIZE),
-    transforms.ToTensor()
-])
-
-transform = transforms.Compose([
-    transforms.Resize(IMAGE_SIZE),
-    transforms.CenterCrop(IMAGE_SIZE),
-    transforms.ToTensor()
-])
-
-dataset = datasets.ImageFolder("data/CUB_200_2011/images")
-
-train_size = int(0.8 * len(dataset))
-val_size = len(dataset) - train_size
-
-train_subset, val_subset = random_split(dataset, [train_size, val_size])
-train_dataset = TransformedSubset(train_subset, transform_augemnt)
-val_dataset = TransformedSubset(val_subset, transform)
-
-train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
-val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
-
-device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
-print("Using device", device)
-
-model = Bird_CNN(c_in=3, c_hidden=16, c_out=200)
-model.to(device)
-optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
-loss_module = nn.CrossEntropyLoss()
-
-trainCNN(model, optimizer, loss_module, train_loader, val_loader, device, 500, SAVE_PATH=SAVE_PATH, save=True)
-exit()
-
-image_path = "test1.jpg"
-
-image = Image.open(image_path).convert("RGB")
-image = transform(image)
-image = image.unsqueeze(0)
-
-confidence, pred = sample(model=model, img=image, device=device, SAVE_PATH=SAVE_PATH)
-print(f"Species: {bird_species(pred.item()).name} | Confidence: {int(confidence.item()*100)/100}")

api/src/common/boxes.py

@@ -0,0 +1,64 @@
+"""Bounding box helpers. Boxes are [xmin, ymin, xmax, ymax] unless noted.
+
+Used by training and by the production server, so keep this free of
+training-only dependencies."""
+
+
+def iou(boxA, boxB):
+    xA = max(boxA[0], boxB[0])
+    yA = max(boxA[1], boxB[1])
+    xB = min(boxA[2], boxB[2])
+    yB = min(boxA[3], boxB[3])
+
+    inter_area = max(0, xB - xA) * max(0, yB - yA)
+
+    boxA_area = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
+    boxB_area = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
+
+    union = boxA_area + boxB_area - inter_area
+
+    return inter_area / union if union > 0 else 0
+
+
+def nms(bboxes, iou_threshold=0.1):
+    """bboxes: [confidence, xmin, ymin, xmax, ymax]. Returns boxes without
+    the confidence, highest confidence first."""
+    bboxes = sorted(bboxes, key=lambda x: x[0], reverse=True)
+
+    keep = []
+    while bboxes:
+        best = bboxes.pop(0)
+        keep.append(best[1:5])
+        bboxes = [box for box in bboxes if iou(best[1:5], box[1:5]) < iou_threshold]
+
+    return keep
+
+
+def xy_center_to_edges(xcenter, ycenter, width, height):
+    width = max(width, 1)
+    height = max(height, 1)
+
+    x = xcenter - (width / 2)
+    y = ycenter - (height / 2)
+
+    return [x, y, x + width, y + height]
+
+
+def cell_to_image_coords(x, y, img_w, img_h, S, cell_i, cell_j):
+    """Grid-cell-relative position -> image position."""
+    return x + (img_w / S) * cell_i, y + (img_h / S) * cell_j
+
+
+def image_to_cell_coords(x, y, img_w, img_h, S, cell_i, cell_j):
+    """Image position -> grid-cell-relative position."""
+    return x - (img_w / S) * cell_i, y - (img_h / S) * cell_j
+
+
+def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
+    gt_cell_i = min(int(x / (img_w / S)), S - 1)
+    gt_cell_j = min(int(y / (img_h / S)), S - 1)
+    return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
+
+
+def scale_box(box, scale_w, scale_h):
+    return [int(box[0] * scale_w), int(box[1] * scale_h), int(box[2] * scale_w), int(box[3] * scale_h)]

api/src/common/checkpoints.py

@@ -0,0 +1,21 @@
+from pathlib import Path
+
+import torch
+
+
+def checkpoint_path(save_dir: str | Path, model_name: str) -> Path:
+    # Layout kept from the original code: <save_dir>/<name>/<name>
+    return Path(save_dir) / model_name / model_name
+
+
+def save_checkpoint(model: torch.nn.Module, save_dir: str | Path, model_name: str) -> Path:
+    path = checkpoint_path(save_dir, model_name)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(model.state_dict(), path)
+    return path
+
+
+def load_checkpoint(model: torch.nn.Module, path: str | Path, device: torch.device) -> torch.nn.Module:
+    state_dict = torch.load(path, map_location=device, weights_only=True)
+    model.load_state_dict(state_dict)
+    return model.to(device).eval()

api/src/common/data.py

@@ -0,0 +1,35 @@
+import torch
+from torch.utils.data import Dataset, random_split
+
+
+class TransformedSubset(Dataset):
+    """Applies a transform to a Subset, so train and validation splits of the
+    same dataset can use different transforms."""
+
+    def __init__(self, subset, transform=None):
+        self.subset = subset
+        self.transform = transform
+
+    def __getitem__(self, idx):
+        x, y = self.subset[idx]
+        if self.transform:
+            x = self.transform(x)
+        return x, y
+
+    def __len__(self):
+        return len(self.subset)
+
+
+def split_dataset(dataset, train_fraction: float, train_transform, val_transform):
+    train_size = int(train_fraction * len(dataset))
+    val_size = len(dataset) - train_size
+    train_subset, val_subset = random_split(dataset, [train_size, val_size])
+    return (
+        TransformedSubset(train_subset, train_transform),
+        TransformedSubset(val_subset, val_transform),
+    )
+
+
+def collate_as_tuples(batch):
+    """For detection datasets whose targets are dicts of different sizes."""
+    return tuple(zip(*batch))

api/src/common/device.py

@@ -0,0 +1,10 @@
+import torch
+
+
+def get_device(name: str = "auto") -> torch.device:
+    """"auto" picks CUDA when available, otherwise CPU."""
+    if name == "auto":
+        name = "cuda:0" if torch.cuda.is_available() else "cpu"
+    device = torch.device(name)
+    print("Using device", device)
+    return device

api/src/common/training.py

@@ -0,0 +1,108 @@
+"""One training loop for all models.
+
+Each epoch trains, validates, prints the metrics and saves the model when the
+monitored metric improves. `prepare_batch` turns a loader batch into
+(inputs, targets); the default just moves both to the device."""
+
+import math
+import time
+
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from src.common.checkpoints import save_checkpoint
+
+
+def _default_prepare(batch, device):
+    inputs, targets = batch
+    return inputs.to(device), targets.to(device)
+
+
+def _run_epoch(model, loss_fn, loader, device, prepare_batch, track_accuracy, optimizer=None, desc=""):
+    training = optimizer is not None
+    model.train(training)
+
+    loss_sum, correct, count = 0.0, 0, 0
+    with torch.set_grad_enabled(training):
+        for batch in tqdm(loader, desc=desc, leave=False):
+            inputs, targets = prepare_batch(batch, device)
+
+            preds = model(inputs)
+            loss = loss_fn(preds, targets)
+
+            if training:
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            loss_sum += loss.item()
+            count += inputs.size(0)
+            if track_accuracy:
+                correct += (preds.argmax(dim=1) == targets).sum().item()
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    metrics = {"loss": loss_sum / max(count, 1)}
+    if track_accuracy:
+        metrics["acc"] = correct / max(count, 1)
+    return metrics
+
+
+def fit(
+    model,
+    loss_fn,
+    optimizer,
+    train_loader,
+    val_loader,
+    device,
+    epochs: int,
+    save_dir: str,
+    model_name: str,
+    save: bool = True,
+    monitor: str = "loss",          # "loss" (lower is better) or "acc" (higher is better)
+    track_accuracy: bool = False,   # for classifiers
+    prepare_batch=_default_prepare,
+):
+    best = math.inf if monitor == "loss" else -math.inf
+
+    for epoch in range(1, epochs + 1):
+        train = _run_epoch(model, loss_fn, train_loader, device, prepare_batch, track_accuracy,
+                           optimizer=optimizer, desc=f"Train {epoch}")
+        val = _run_epoch(model, loss_fn, val_loader, device, prepare_batch, track_accuracy,
+                         desc=f"Validate {epoch}")
+
+        improved = val[monitor] < best if monitor == "loss" else val[monitor] > best
+        if improved:
+            best = val[monitor]
+            if save:
+                save_checkpoint(model, save_dir, model_name)
+
+        # Loss is the summed batch loss per sample, shown x1000.
+        parts = [f"epoch: {epoch}"]
+        if track_accuracy:
+            parts.append(f"train acc: {train['acc'] * 100:.1f}% | val acc: {val['acc'] * 100:.1f}%")
+        parts.append(f"train loss: {train['loss'] * 1000:.2f} | val loss: {val['loss'] * 1000:.2f}")
+        if improved and save:
+            parts.append("saved")
+        print(" | ".join(parts))
+
+    return best
+
+
+def benchmark_num_workers(dataset, model, device, batch_size=64, rounds=2):
+    """Prints how long `rounds` passes over the dataset take for different
+    DataLoader num_workers values, to pick the fastest one."""
+    import multiprocessing as mp
+
+    model.eval()
+    for num_workers in range(0, mp.cpu_count(), 2):
+        loader = DataLoader(dataset, shuffle=True, num_workers=num_workers,
+                            batch_size=batch_size, pin_memory=True)
+        start = time.time()
+        with torch.no_grad():
+            for _ in range(rounds):
+                for images, _ in tqdm(loader, leave=False):
+                    model(images.to(device))
+        print(f"num_workers={num_workers}: {time.time() - start:.0f} seconds")

api/src/common/visualize.py

@@ -0,0 +1,26 @@
+"""Plot helpers for local experiments (matplotlib is not installed in the
+production image, so it is imported lazily)."""
+
+
+def show_image(image, title=None):
+    """image: tensor [C, H, W]."""
+    from matplotlib import pyplot as plt
+
+    plt.figure(figsize=(3, 3))
+    plt.imshow(image.permute(1, 2, 0))
+    if title:
+        plt.title(title)
+    plt.axis("off")
+    plt.show()
+
+
+def show_batch(loader, count=4):
+    from matplotlib import pyplot as plt
+
+    images, labels = next(iter(loader))
+    for i in range(min(count, len(images))):
+        plt.figure(figsize=(3, 3))
+        plt.imshow(images[i].permute(1, 2, 0))
+        plt.title(f"Label: {labels[i].item()}")
+        plt.axis("off")
+    plt.show()

api/src/main.py

@@ -1,8 +1,96 @@
-from src.yolo.train_yolo_faces import train_yolo
+"""One entrypoint for everything in api/. Run it from the api/ folder:
+
+    python -m src.main                              list all tasks
+    python -m src.main yolo.train                   run a task with its defaults
+    python -m src.main yolo.train --epochs 20 --batch-size 32
+    python -m src.main yolo.train --help            show every option of a task
+
+Defaults live in src/<group>/config.py; edit them there or override them per
+run with flags. To start from an IDE "Run" button without arguments, set
+DEFAULT_TASK (and optionally DEFAULT_ARGS) below.
+"""
+
+import argparse
+from dataclasses import fields
+import importlib
+import sys
+
+# Used when the script is started without arguments.
+DEFAULT_TASK = None            # e.g. "yolo.train"
+DEFAULT_ARGS = []              # e.g. ["--epochs", "20"]
+
+# group name -> package that contains config.py (Config) and tasks.py (TASKS).
+# Packages are only imported when one of their tasks is used, so a missing
+# training dependency does not break the others.
+GROUPS = {
+    "bird": "src.bird_cnn",
+    "yolo": "src.yolo",
+    "rcnn": "src.r_cnn",
+    "server": "src.production",
+}
 
 
-def main():
+def load_group(group: str):
-    train_yolo()
+    package = GROUPS[group]
+    config_cls = importlib.import_module(f"{package}.config").Config
+    tasks = importlib.import_module(f"{package}.tasks").TASKS
+    return config_cls, tasks
+
+
+def first_line(text):
+    return (text or "").strip().splitlines()[0] if text and text.strip() else ""
+
+
+def list_tasks():
+    print(__doc__)
+    print("Tasks:")
+    for group in GROUPS:
+        try:
+            _, tasks = load_group(group)
+        except ImportError as e:
+            print(f"  {group}.*  (unavailable: {e})")
+            continue
+        for name, fn in tasks.items():
+            print(f"  {group + '.' + name:<28} {first_line(fn.__doc__)}")
+
+
+def build_parser(task: str, config_cls, fn):
+    parser = argparse.ArgumentParser(
+        prog=f"python -m src.main {task}",
+        description=(fn.__doc__ or "").strip(),
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    for f in fields(config_cls):
+        flag = "--" + f.name.replace("_", "-")
+        if f.type is bool:
+            parser.add_argument(flag, dest=f.name, action=argparse.BooleanOptionalAction, default=f.default)
+        else:
+            parser.add_argument(flag, dest=f.name, type=f.type, default=f.default, metavar=f.type.__name__.upper())
+    return parser
+
+
+def main(argv=None):
+    argv = sys.argv[1:] if argv is None else argv
+    if not argv and DEFAULT_TASK:
+        argv = [DEFAULT_TASK, *DEFAULT_ARGS]
+    if not argv or argv[0] in ("-h", "--help", "list"):
+        list_tasks()
+        return
+
+    task, *rest = argv
+    group, _, name = task.partition(".")
+    if group not in GROUPS:
+        sys.exit(f"Unknown group '{group}'. Groups: {', '.join(GROUPS)}")
+
+    config_cls, tasks = load_group(group)
+    if name not in tasks:
+        sys.exit(f"Unknown task '{task}'. Tasks in {group}: {', '.join(tasks)}")
+
+    fn = tasks[name]
+    args = build_parser(task, config_cls, fn).parse_args(rest)
+    cfg = config_cls(**vars(args))
+    fn(cfg)
+
 
 if __name__ == "__main__":
     main()

api/src/models/bird_cnn.py

@@ -1,24 +1,22 @@
+"""Bird classifier. Shared by training and the production server; the layer
+names must not change or saved weights stop loading."""
+
+from enum import Enum
+
 import torch.nn as nn
 import torch.nn.functional as F
 
-class  SeparableConvolution(nn.Module):
-    def __init__(self, c_in, c_out, kernel_size):
-        super().__init__()
-        self.depthwise = nn.Conv2d(c_in, c_in, kernel_size, groups=c_in, padding=kernel_size//2)
-        self.bn1 = nn.BatchNorm2d(c_in)
-        self.pointwise = nn.Conv2d(c_in, c_out, kernel_size=1)
-        self.bn2 = nn.BatchNorm2d(c_out)
 
-    def forward(self, x):
+class BirdSpecies(Enum):
-        x = self.depthwise(x)
+    """Classes of the deployed 7-species model (build_models/bird_cnn)."""
-        x = self.bn1(x)
+    Common_Kingfisher = 0
-        x = F.relu(x)
+    Common_Myna = 1
+    House_Crow = 2
+    Indian_Peacock = 3
+    Indian_Pitta = 4
+    Ruddy_Shelduck = 5
+    Sarus_Crane = 6
 
-        x = self.pointwise(x)
-        x = self.bn2(x)
-        x = F.relu(x)
-
-        return x
 
 class SkipBlock(nn.Module):
     def __init__(self, c_in, c_out, kernel_size=3):
@@ -41,7 +39,7 @@ class SkipBlock(nn.Module):
         )
 
     def forward(self, x):
-        return(F.relu(self.conv_skip(x) + self.conv(x), inplace=True))
+        return F.relu(self.conv_skip(x) + self.conv(x), inplace=True)
 
 
 class Bird_CNN(nn.Module):

api/src/models/r_cnn.py

@@ -0,0 +1,36 @@
+import torch.nn as nn
+
+
+class ObjectDetectionCNN(nn.Module):
+    """Small crop classifier used by the R-CNN style experiments
+    (selective search proposals -> classify each crop)."""
+
+    def __init__(self, c_in, c_hidden, c_out, layers):
+        super().__init__()
+
+        self.model = nn.ModuleList()
+
+        self.model.append(nn.Sequential(
+            nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
+            nn.BatchNorm2d(c_hidden),
+            nn.ReLU(inplace=True)
+        ))
+
+        for _ in range(layers - 1):
+            self.model.append(nn.Sequential(
+                nn.Conv2d(c_hidden, c_hidden, kernel_size=3, padding=1),
+                nn.BatchNorm2d(c_hidden),
+                nn.ReLU(inplace=True)
+            ))
+
+        self.model.append(nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Flatten(),
+            nn.Linear(c_hidden, c_out),
+            nn.Dropout(0.3)
+        ))
+
+    def forward(self, x):
+        for layer in self.model:
+            x = layer(x)
+        return x

api/src/models/yolo.py

@@ -0,0 +1,128 @@
+"""YOLO-style single-shot face detector. Shared by training and the
+production server; the layer names must not change or saved weights stop
+loading."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from src.common.boxes import cell_to_image_coords, nms, xy_center_to_edges
+
+
+class SkipBlock(nn.Module):
+    def __init__(self, c_in, c_out, kernel_size=3):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size // 2),
+            nn.GroupNorm(num_groups=c_out // 8, num_channels=c_out),
+            nn.LeakyReLU(inplace=True),
+
+            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size // 2),
+            nn.GroupNorm(num_groups=c_out // 8, num_channels=c_out),
+            nn.LeakyReLU(inplace=True),
+
+            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size // 2),
+            nn.GroupNorm(num_groups=c_out // 8, num_channels=c_out),
+            nn.LeakyReLU(inplace=True)
+        )
+        self.conv_skip = nn.Sequential(
+            nn.Conv2d(c_in, c_out, 1),
+            nn.GroupNorm(num_groups=c_out // 8, num_channels=c_out),
+            nn.LeakyReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        # F.dropout defaults to training=True, so this stays active in eval
+        # mode too. Kept as-is because the saved models were trained this way.
+        return F.dropout(F.leaky_relu(self.conv_skip(x) + self.conv(x), inplace=True), p=0.3)
+
+
+class Yolo_model(nn.Module):
+    """Output: [batch, grid, grid, boxes*5 + labels] with per cell
+    (x, y, w, h, confidence, classes...).
+
+    size_activation decides how w/h are produced:
+      "exp"     - exp(x), used by the deployed build_models/face_detection_yolo
+      "sigmoid" - sigmoid(x), what the training code switched to later
+    """
+
+    def __init__(self, c_in, boxes, grid, labels, c_hidden=16, size_activation="sigmoid"):
+        super().__init__()
+        if size_activation not in ("exp", "sigmoid"):
+            raise ValueError(f"size_activation must be 'exp' or 'sigmoid', got {size_activation!r}")
+        self.size_activation = size_activation
+
+        self.model = nn.Sequential(
+            nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
+            nn.GroupNorm(num_groups=c_hidden // 8, num_channels=c_hidden),
+            nn.LeakyReLU(inplace=True),
+
+            SkipBlock(c_in=c_hidden, c_out=c_hidden),
+            SkipBlock(c_in=c_hidden, c_out=c_hidden),
+            SkipBlock(c_in=c_hidden, c_out=c_hidden),
+            SkipBlock(c_in=c_hidden, c_out=c_hidden),
+
+            SkipBlock(c_in=c_hidden, c_out=c_hidden * 2),
+            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
+            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
+            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
+
+            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 4),
+            SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
+            SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
+            SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
+
+            nn.Conv2d(c_hidden * 4, c_hidden * 8, kernel_size=3, padding=1),
+            nn.GroupNorm(num_groups=c_hidden // 2, num_channels=c_hidden * 8),
+            nn.LeakyReLU(inplace=True),
+            nn.Dropout(0.3),
+
+            nn.AdaptiveAvgPool2d((grid, grid)),
+            nn.Conv2d(c_hidden * 8, boxes * 5 + labels, kernel_size=1)
+        )
+
+    def forward(self, x):
+        x = self.model(x).permute(0, 2, 3, 1)
+        if self.size_activation == "sigmoid":
+            return F.sigmoid(x)
+        center = F.sigmoid(x[..., :2])
+        size = torch.exp(x[..., 2:4])
+        conf_class = F.sigmoid(x[..., 4:])
+        return torch.cat([center, size, conf_class], dim=3)
+
+
+def convert_prediction(label, image, threshold=0.9):
+    """Turns one model output (or target) [grid, grid, 5 + classes] into boxes
+    in image pixels. Returns (boxes after NMS, cells with an object, cells
+    without one); the cell lists are only used for visualisation."""
+    image = image.clone().detach()
+    label = label.clone().detach()
+
+    image_size = image.shape[1]
+    grid_number = label.shape[0]
+    grid_size = image_size / grid_number
+
+    boxes_to_draw = []
+    grids_to_draw_obj = []
+    grids_to_draw_noobj = []
+    for x in range(label.shape[0]):
+        for y in range(label.shape[1]):
+            cell = [x * grid_size, y * grid_size, (x + 1) * grid_size, (y + 1) * grid_size]
+            if label[x, y, 4].item() < threshold:
+                grids_to_draw_noobj.append(cell)
+                continue
+            grids_to_draw_obj.append(cell)
+
+            boxx = label[x, y, 0] * (image_size / grid_number)
+            boxy = label[x, y, 1] * (image_size / grid_number)
+
+            boxw = label[x, y, 2] * image_size
+            boxh = label[x, y, 3] * image_size
+
+            boxx, boxy = cell_to_image_coords(x=boxx, y=boxy, img_w=image_size, img_h=image_size,
+                                              S=grid_number, cell_i=x, cell_j=y)
+
+            boxes_to_draw.append([label[x, y, 4]] + xy_center_to_edges(boxx, boxy, boxw, boxh))
+
+    boxes_to_draw = nms(boxes_to_draw)
+    return boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj

api/src/production/app.py

@@ -0,0 +1,41 @@
+"""FastAPI app. Start with: uvicorn src.production.app:app (from api/),
+or: python -m src.main server.run"""
+
+from contextlib import asynccontextmanager
+import logging
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from src.production import predict, reviews, settings
+from src.production.inference import load_models
+
+logger = logging.getLogger("server")
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    logging.basicConfig(level=logging.INFO)
+    if not settings.API_KEY:
+        logger.warning("API_KEY is not set: GET/DELETE /review are disabled")
+    load_models()
+    reviews.init_db()
+    yield
+
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=settings.ORIGINS,
+    allow_methods=["GET", "POST", "DELETE"],
+    allow_headers=["Authorization", "Content-Type"],
+)
+
+app.include_router(predict.router)
+app.include_router(reviews.router)
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}

api/src/production/config.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Config:
+    """How to start the API server locally. Server limits (rate limits,
+    parallel inferences, ...) are in src/production/settings.py and are set
+    with environment variables."""
+
+    host: str = "127.0.0.1"
+    port: int = 8000
+    reload: bool = True
+    ws_max_size: int = 1024 * 1024   # largest websocket message in bytes

api/src/production/inference.py

@@ -0,0 +1,85 @@
+"""Loading the models and running one prediction. These functions are
+blocking and run on the InferenceGate worker threads."""
+
+import io
+
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+
+from src.common.boxes import scale_box
+from src.common.checkpoints import load_checkpoint
+from src.models.bird_cnn import BirdSpecies, Bird_CNN
+from src.models.yolo import Yolo_model, convert_prediction
+from src.production import settings
+
+Image.MAX_IMAGE_PIXELS = settings.MAX_IMAGE_PIXELS
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+models = {}
+
+transform_bird = transforms.Compose([
+    transforms.Resize(settings.BIRD_IMAGE_SIZE),
+    transforms.CenterCrop(settings.BIRD_IMAGE_SIZE),
+    transforms.ToTensor()
+])
+
+transform_face = transforms.Compose([
+    transforms.Resize((settings.FACE_IMAGE_SIZE, settings.FACE_IMAGE_SIZE)),
+    transforms.ToTensor()
+])
+
+
+def load_models():
+    bird = Bird_CNN(c_in=3, c_hidden=settings.BIRD_C_HIDDEN, c_out=len(BirdSpecies))
+    models["bird"] = load_checkpoint(bird, settings.BUILD_PATH / settings.BIRD_MODEL_FILE, device)
+
+    face = Yolo_model(c_in=3, boxes=1, grid=settings.FACE_GRID, labels=1, c_hidden=settings.FACE_C_HIDDEN,
+                      size_activation=settings.FACE_SIZE_ACTIVATION)
+    models["face"] = load_checkpoint(face, settings.BUILD_PATH / settings.FACE_MODEL_FILE, device)
+
+
+class InvalidImage(Exception):
+    pass
+
+
+def open_image(data: bytes) -> tuple[Image.Image, tuple[int, int]]:
+    """Open lazily, reject oversized images before decoding, decode JPEGs at
+    reduced scale. Returns the image and its ORIGINAL (width, height)."""
+    try:
+        image = Image.open(io.BytesIO(data))
+        original_size = image.size
+        if original_size[0] * original_size[1] > settings.MAX_IMAGE_PIXELS:
+            raise InvalidImage("Image too large")
+        image.draft("RGB", settings.DECODE_DRAFT_SIZE)
+        return image.convert("RGB"), original_size
+    except InvalidImage:
+        raise
+    except Exception as e:  # PIL raises many different types for bad data
+        raise InvalidImage("Not a valid image") from e
+
+
+def classify_bird(data: bytes) -> dict:
+    image, _ = open_image(data)
+    tensor = transform_bird(image).unsqueeze(0).to(device)
+    with torch.inference_mode():
+        probs = F.softmax(models["bird"](tensor), dim=1)
+        confidence, cls = torch.max(probs, dim=1)
+    return {
+        "class": BirdSpecies(cls.item()).name,
+        "confidence": confidence.item()
+    }
+
+
+def detect_faces(data: bytes) -> list:
+    image, (W, H) = open_image(data)
+    scale_w = W / settings.FACE_IMAGE_SIZE
+    scale_h = H / settings.FACE_IMAGE_SIZE
+
+    tensor = transform_face(image).unsqueeze(0).to(device)
+    with torch.inference_mode():
+        pred = models["face"](tensor)
+        bboxes, _, _ = convert_prediction(pred.squeeze(0), tensor.squeeze(0), threshold=settings.FACE_THRESHOLD)
+
+    return [scale_box(b, scale_w, scale_h) for b in bboxes]

api/src/production/limits.py

@@ -0,0 +1,95 @@
+"""Overload protection: per-IP rate limits and the inference thread pool."""
+
+import asyncio
+from collections import defaultdict, deque
+from concurrent.futures import ThreadPoolExecutor
+import time
+
+from fastapi import HTTPException, Request, WebSocket
+import torch
+
+from src.production import settings
+
+
+def client_ip(conn: Request | WebSocket) -> str:
+    # Caddy is the only proxy in front of this container and appends the real
+    # client address as the last X-Forwarded-For entry.
+    forwarded = conn.headers.get("x-forwarded-for")
+    if forwarded:
+        return forwarded.split(",")[-1].strip()
+    return conn.client.host if conn.client else "unknown"
+
+
+class RateLimiter:
+    """Sliding-window limiter per key (in memory, single process)."""
+
+    def __init__(self, limit: int, window: float):
+        self.limit = limit
+        self.window = window
+        self.hits = defaultdict(deque)
+
+    def allow(self, key: str) -> bool:
+        now = time.monotonic()
+        if len(self.hits) > 10_000:
+            self.hits = defaultdict(deque, {
+                k: v for k, v in self.hits.items() if v and now - v[-1] < self.window
+            })
+        q = self.hits[key]
+        while q and now - q[0] >= self.window:
+            q.popleft()
+        if len(q) >= self.limit:
+            return False
+        q.append(now)
+        return True
+
+
+predict_limiter = RateLimiter(settings.PREDICT_RATE_LIMIT, settings.PREDICT_RATE_WINDOW)
+review_limiter = RateLimiter(settings.REVIEW_RATE_LIMIT, settings.REVIEW_RATE_WINDOW)
+
+
+def rate_limit(limiter: RateLimiter):
+    """FastAPI dependency that answers 429 when the client is over the limit."""
+    def dependency(request: Request):
+        if not limiter.allow(client_ip(request)):
+            raise HTTPException(status_code=429, detail="Too many requests", headers={"Retry-After": "60"})
+    return dependency
+
+
+class Busy(Exception):
+    pass
+
+
+def init_inference_thread():
+    # Must run inside each worker thread: with OpenMP the thread count is a
+    # per-thread setting, so setting it once in the main thread is not enough.
+    torch.set_num_threads(settings.TORCH_THREADS)
+
+
+class InferenceGate:
+    """Runs predictions on `parallel` worker threads. At most
+    `parallel + queued` jobs (running + waiting) are admitted; everything
+    else is rejected at once instead of queueing up and eating memory.
+
+    The models are in eval mode under torch.inference_mode(), so several
+    threads can safely run forward passes on the same model object."""
+
+    def __init__(self, parallel: int, queued: int):
+        self.max_pending = parallel + queued
+        self.pending = 0
+        self.executor = ThreadPoolExecutor(
+            max_workers=parallel,
+            thread_name_prefix="inference",
+            initializer=init_inference_thread,
+        )
+
+    async def run(self, fn, *args):
+        if self.pending >= self.max_pending:
+            raise Busy
+        self.pending += 1
+        try:
+            return await asyncio.get_running_loop().run_in_executor(self.executor, fn, *args)
+        finally:
+            self.pending -= 1
+
+
+gate = InferenceGate(settings.PARALLEL_INFERENCES, settings.QUEUED_INFERENCES)

api/src/production/predict.py

@@ -0,0 +1,97 @@
+"""Routes: bird classification (POST /predict) and face detection
+(websocket /predict_face)."""
+
+from collections import defaultdict
+import logging
+import time
+
+from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
+
+from src.production import settings
+from src.production.inference import InvalidImage, classify_bird, detect_faces
+from src.production.limits import Busy, client_ip, gate, predict_limiter, rate_limit
+
+logger = logging.getLogger("server")
+router = APIRouter()
+
+
+async def read_limited(file: UploadFile, limit: int) -> bytes:
+    data = bytearray()
+    while chunk := await file.read(64 * 1024):
+        data += chunk
+        if len(data) > limit:
+            raise HTTPException(status_code=413, detail="File too large")
+    return bytes(data)
+
+
+@router.post("/predict", dependencies=[Depends(rate_limit(predict_limiter))])
+async def predict(file: UploadFile = File(...)):
+    data = await read_limited(file, settings.MAX_UPLOAD_BYTES)
+    try:
+        return await gate.run(classify_bird, data)
+    except Busy:
+        raise HTTPException(status_code=503, detail="Server busy, try again shortly", headers={"Retry-After": "5"})
+    except InvalidImage as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+active_websockets = defaultdict(int)
+
+
+@router.websocket("/predict_face")
+async def predict_face(websocket: WebSocket):
+    ip = client_ip(websocket)
+
+    # CORS does not apply to websockets, so check the origin ourselves.
+    if websocket.headers.get("origin") not in settings.ORIGINS:
+        await websocket.close(code=1008)
+        return
+    if (sum(active_websockets.values()) >= settings.MAX_WEBSOCKETS
+            or active_websockets[ip] >= settings.MAX_WEBSOCKETS_PER_IP):
+        await websocket.close(code=1013)  # try again later
+        return
+
+    active_websockets[ip] += 1
+    try:
+        await websocket.accept()
+        last = 0.0
+
+        while True:
+            message = await websocket.receive()
+            if message["type"] == "websocket.disconnect":
+                break
+
+            data = message.get("bytes")
+            if data is None or len(data) > settings.MAX_FRAME_BYTES:
+                await websocket.close(code=1009 if data else 1003)
+                break
+
+            # Drop frames that arrive too fast or while the server is busy
+            # instead of queueing them.
+            now = time.monotonic()
+            if now - last < settings.FACE_MIN_INTERVAL:
+                continue
+            last = now
+
+            try:
+                boxes = await gate.run(detect_faces, data)
+            except Busy:
+                continue
+            except InvalidImage:
+                await websocket.close(code=1003)
+                break
+
+            await websocket.send_json({"bboxes": boxes})
+
+    except WebSocketDisconnect:
+        pass
+    except Exception:
+        logger.exception("predict_face failed")
+        try:
+            await websocket.close(code=1011)
+        except Exception:
+            pass
+    finally:
+        active_websockets[ip] -= 1
+        if active_websockets[ip] <= 0:
+            del active_websockets[ip]

api/src/production/reviews.py

@@ -0,0 +1,84 @@
+"""Routes and storage for reviews (SQLite)."""
+
+from contextlib import closing
+from datetime import datetime, timezone
+import secrets
+import sqlite3
+
+from fastapi import APIRouter, Depends, Header, HTTPException, Query
+from pydantic import BaseModel, Field
+
+from src.production import settings
+from src.production.limits import rate_limit, review_limiter
+
+router = APIRouter()
+
+
+def init_db():
+    settings.DATABASE_DIR.mkdir(parents=True, exist_ok=True)
+    with closing(sqlite3.connect(settings.DATABASE)) as conn:
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS reviews (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                website TEXT NOT NULL,
+                rating INTEGER NOT NULL,
+                text TEXT NOT NULL,
+                created_at TEXT NOT NULL
+            )
+        """)
+        conn.commit()
+
+
+def connect_db() -> sqlite3.Connection:
+    conn = sqlite3.connect(settings.DATABASE, timeout=5)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def get_api_key(authorization: str = Header(None)):
+    if not settings.API_KEY or not authorization or not secrets.compare_digest(
+        authorization.encode(), f"Bearer {settings.API_KEY}".encode()
+    ):
+        raise HTTPException(status_code=401, detail="Unauthorized")
+    return authorization
+
+
+class Review(BaseModel):
+    website: str = Field(min_length=1, max_length=100)
+    rating: int = Field(ge=1, le=5)
+    text: str = Field(max_length=1000)
+    # The client also sends a "date"; it is ignored, the server sets the time.
+
+
+@router.get("/review")
+def get_reviews(
+    limit: int = Query(100, ge=1, le=500),
+    offset: int = Query(0, ge=0),
+    auth=Depends(get_api_key),
+):
+    with closing(connect_db()) as conn:
+        rows = conn.execute(
+            "SELECT * FROM reviews ORDER BY id DESC LIMIT ? OFFSET ?", (limit, offset)
+        ).fetchall()
+    return {"data": [dict(row) for row in rows]}
+
+
+# Public on purpose: no credentials needed, only rate limited per IP.
+@router.post("/review", dependencies=[Depends(rate_limit(review_limiter))])
+def post_review(review: Review):
+    created_at = datetime.now(timezone.utc).isoformat()
+    with closing(connect_db()) as conn:
+        conn.execute(
+            "INSERT INTO reviews (website, rating, text, created_at) VALUES (?, ?, ?, ?)",
+            (review.website, review.rating, review.text.strip(), created_at),
+        )
+        conn.commit()
+    return {"status": "ok"}
+
+
+@router.delete("/review")
+def delete_review(auth=Depends(get_api_key)):
+    with closing(connect_db()) as conn:
+        conn.execute("DELETE FROM reviews")
+        conn.commit()
+    return {"status": "ok"}

api/src/production/server.py

@@ -1,429 +0,0 @@
-import asyncio
-from collections import defaultdict, deque
-from concurrent.futures import ThreadPoolExecutor
-from contextlib import asynccontextmanager, closing
-from datetime import datetime, timezone
-from enum import Enum
-import io
-import logging
-import os
-from pathlib import Path
-import secrets
-import sqlite3
-import time
-
-from dotenv import load_dotenv
-from fastapi import Depends, FastAPI, File, Header, HTTPException, Query, Request, UploadFile, WebSocket, WebSocketDisconnect
-from fastapi.middleware.cors import CORSMiddleware
-from PIL import Image
-from pydantic import BaseModel, Field
-import torch
-import torch.nn.functional as F
-from torchvision import transforms
-
-from yolo_model_production import convert_prediction, Yolo_model
-from bird_cnn_production import Bird_CNN
-
-logger = logging.getLogger("server")
-
-BASE_DIR = Path(__file__).resolve().parent
-BUILD_PATH = Path(os.getenv("BUILD_PATH", BASE_DIR / "build_models"))
-DATABASE_DIR = Path(os.getenv("DATABASE_DIR", BASE_DIR / "database"))
-DATABASE = DATABASE_DIR / "reviews.db"
-
-load_dotenv(DATABASE_DIR / ".env")
-API_KEY = os.getenv("API_KEY")  # if missing, protected endpoints reject every request
-
-IMAGE_SIZE_CNN = 64
-IMAGE_SIZE_YOLO = 64
-
-# ---- Resource limits (the server is weak, keep everything small and bounded) ----
-# Number of predictions (bird or face) that may run at the same time.
-PARALLEL_INFERENCES = max(1, int(os.getenv("PARALLEL_INFERENCES", "1")))
-# Extra requests allowed to wait for a free slot; anything beyond is rejected.
-QUEUED_INFERENCES = max(0, int(os.getenv("QUEUED_INFERENCES", "1")))
-# PyTorch threads used by EACH running prediction. Total CPU use is roughly
-# PARALLEL_INFERENCES * TORCH_THREADS, so keep the product <= your CPU cores.
-TORCH_THREADS = max(1, int(os.getenv("TORCH_THREADS", "1")))
-MAX_UPLOAD_BYTES = int(max(1, float(os.getenv("MAX_UPLOAD_MEGABYTES", "15"))) * 1024 * 1024)
-MAX_FRAME_BYTES = 1 * 1024 * 1024
-MAX_IMAGE_PIXELS = 20_000_000
-MAX_WEBSOCKETS = 3
-MAX_WEBSOCKETS_PER_IP = 1
-FACE_MAX_FPS = float(os.getenv("FACE_MAX_FPS", "5.5"))  # per websocket connection
-FACE_MIN_INTERVAL = 1 / FACE_MAX_FPS
-DECODE_DRAFT_SIZE = (256, 256)  # JPEG decodes at reduced scale, still larger than the 64px model input
-
-Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
-
-origins = [
-    "https://marvinkrausser.com",
-]
-
-
-class BirdSpecies(Enum):
-    Common_Kingfisher = 0
-    Common_Myna = 1
-    House_Crow = 2
-    Indian_Peacock = 3
-    Indian_Pitta = 4
-    Ruddy_Shelduck = 5
-    Sarus_Crane = 6
-
-
-transform_bird = transforms.Compose([
-    transforms.Resize(IMAGE_SIZE_CNN),
-    transforms.CenterCrop(IMAGE_SIZE_CNN),
-    transforms.ToTensor()
-])
-
-transform_face = transforms.Compose([
-    transforms.Resize((IMAGE_SIZE_YOLO, IMAGE_SIZE_YOLO)),
-    transforms.ToTensor()
-])
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-models = {}
-
-
-def load_models():
-    model_bird = Bird_CNN(c_in=3, c_hidden=16, c_out=7)
-    model_bird.load_state_dict(torch.load(BUILD_PATH / "bird_cnn", map_location=device, weights_only=True))
-    model_bird.to(device).eval()
-
-    model_face = Yolo_model(c_in=3, boxes=1, grid=6, labels=1, c_hidden=16)
-    model_face.load_state_dict(torch.load(BUILD_PATH / "face_detection_yolo", map_location=device, weights_only=True))
-    model_face.to(device).eval()
-
-    models["bird"] = model_bird
-    models["face"] = model_face
-
-
-def init_db():
-    DATABASE_DIR.mkdir(parents=True, exist_ok=True)
-    with closing(sqlite3.connect(DATABASE)) as conn:
-        conn.execute("""
-            CREATE TABLE IF NOT EXISTS reviews (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                website TEXT NOT NULL,
-                rating INTEGER NOT NULL,
-                text TEXT NOT NULL,
-                created_at TEXT NOT NULL
-            )
-        """)
-        conn.commit()
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    logging.basicConfig(level=logging.INFO)
-    if not API_KEY:
-        logger.warning("API_KEY is not set: GET/DELETE /review are disabled")
-    load_models()
-    init_db()
-    yield
-
-
-app = FastAPI(lifespan=lifespan)
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_methods=["GET", "POST", "DELETE"],
-    allow_headers=["Authorization", "Content-Type"],
-)
-
-
-# --------------------------------------------------------------------------
-# Overload protection
-# --------------------------------------------------------------------------
-
-def client_ip(conn: Request | WebSocket) -> str:
-    # Caddy is the only proxy in front of this container and appends the real
-    # client address as the last X-Forwarded-For entry.
-    forwarded = conn.headers.get("x-forwarded-for")
-    if forwarded:
-        return forwarded.split(",")[-1].strip()
-    return conn.client.host if conn.client else "unknown"
-
-
-class RateLimiter:
-    """Sliding-window limiter per key (in memory, single process)."""
-
-    def __init__(self, limit: int, window: float):
-        self.limit = limit
-        self.window = window
-        self.hits = defaultdict(deque)
-
-    def allow(self, key: str) -> bool:
-        now = time.monotonic()
-        if len(self.hits) > 10_000:
-            self.hits = defaultdict(deque, {
-                k: v for k, v in self.hits.items() if v and now - v[-1] < self.window
-            })
-        q = self.hits[key]
-        while q and now - q[0] >= self.window:
-            q.popleft()
-        if len(q) >= self.limit:
-            return False
-        q.append(now)
-        return True
-
-
-predict_limiter = RateLimiter(
-    limit=max(1, int(os.getenv("PREDICT_RATE_LIMIT", "10"))),
-    window=max(1.0, float(os.getenv("PREDICT_RATE_WINDOW", "60"))),
-)
-review_limiter = RateLimiter(
-    limit=max(1, int(os.getenv("REVIEW_RATE_LIMIT", "5"))),
-    window=max(1.0, float(os.getenv("REVIEW_RATE_WINDOW", "60"))),
-)
-
-
-def rate_limit(limiter: RateLimiter):
-    def dependency(request: Request):
-        if not limiter.allow(client_ip(request)):
-            raise HTTPException(status_code=429, detail="Too many requests", headers={"Retry-After": "60"})
-    return dependency
-
-
-class Busy(Exception):
-    pass
-
-
-def init_inference_thread():
-    # Must run inside each worker thread: with OpenMP the thread count is a
-    # per-thread setting, so setting it once in the main thread is not enough.
-    torch.set_num_threads(TORCH_THREADS)
-
-
-class InferenceGate:
-    """Runs predictions on `parallel` worker threads. At most
-    `parallel + queued` jobs (running + waiting) are admitted; everything
-    else is rejected at once instead of queueing up and eating memory.
-
-    The models are in eval mode under torch.inference_mode(), so several
-    threads can safely run forward passes on the same model object."""
-
-    def __init__(self, parallel: int, queued: int):
-        self.max_pending = parallel + queued
-        self.pending = 0
-        self.executor = ThreadPoolExecutor(
-            max_workers=parallel,
-            thread_name_prefix="inference",
-            initializer=init_inference_thread,
-        )
-
-    async def run(self, fn, *args):
-        if self.pending >= self.max_pending:
-            raise Busy
-        self.pending += 1
-        try:
-            return await asyncio.get_running_loop().run_in_executor(self.executor, fn, *args)
-        finally:
-            self.pending -= 1
-
-
-gate = InferenceGate(PARALLEL_INFERENCES, QUEUED_INFERENCES)
-
-
-class InvalidImage(Exception):
-    pass
-
-
-def open_image(data: bytes) -> tuple[Image.Image, tuple[int, int]]:
-    """Open lazily, reject oversized images before decoding, decode JPEGs at
-    reduced scale. Returns the image and its ORIGINAL (width, height)."""
-    try:
-        image = Image.open(io.BytesIO(data))
-        original_size = image.size
-        if original_size[0] * original_size[1] > MAX_IMAGE_PIXELS:
-            raise InvalidImage("Image too large")
-        image.draft("RGB", DECODE_DRAFT_SIZE)
-        return image.convert("RGB"), original_size
-    except InvalidImage:
-        raise
-    except Exception as e:  # PIL raises many different types for bad data
-        raise InvalidImage("Not a valid image") from e
-
-
-# --------------------------------------------------------------------------
-# Bird classification
-# --------------------------------------------------------------------------
-
-def classify_bird(data: bytes) -> dict:
-    image, _ = open_image(data)
-    tensor = transform_bird(image).unsqueeze(0).to(device)
-    with torch.inference_mode():
-        probs = F.softmax(models["bird"](tensor), dim=1)
-        confidence, cls = torch.max(probs, dim=1)
-    return {
-        "class": BirdSpecies(cls.item()).name,
-        "confidence": confidence.item()
-    }
-
-
-async def read_limited(file: UploadFile, limit: int) -> bytes:
-    data = bytearray()
-    while chunk := await file.read(64 * 1024):
-        data += chunk
-        if len(data) > limit:
-            raise HTTPException(status_code=413, detail="File too large")
-    return bytes(data)
-
-
-@app.post("/predict", dependencies=[Depends(rate_limit(predict_limiter))])
-async def predict(file: UploadFile = File(...)):
-    data = await read_limited(file, MAX_UPLOAD_BYTES)
-    try:
-        return await gate.run(classify_bird, data)
-    except Busy:
-        raise HTTPException(status_code=503, detail="Server busy, try again shortly", headers={"Retry-After": "5"})
-    except InvalidImage as e:
-        raise HTTPException(status_code=400, detail=str(e))
-
-
-# --------------------------------------------------------------------------
-# Face detection (websocket)
-# --------------------------------------------------------------------------
-
-def detect_faces(data: bytes) -> list:
-    image, (W, H) = open_image(data)
-    scale_w = W / IMAGE_SIZE_YOLO
-    scale_h = H / IMAGE_SIZE_YOLO
-
-    tensor = transform_face(image).unsqueeze(0).to(device)
-    with torch.inference_mode():
-        pred = models["face"](tensor)
-        bboxes, _, _ = convert_prediction(pred.squeeze(0), tensor.squeeze(0), threshold=0.9)
-
-    return [
-        [int(b[0] * scale_w), int(b[1] * scale_h), int(b[2] * scale_w), int(b[3] * scale_h)]
-        for b in bboxes
-    ]
-
-
-active_websockets = defaultdict(int)
-
-
-@app.websocket("/predict_face")
-async def predict_face(websocket: WebSocket):
-    ip = client_ip(websocket)
-
-    # CORS does not apply to websockets, so check the origin ourselves.
-    if websocket.headers.get("origin") not in origins:
-        await websocket.close(code=1008)
-        return
-    if (sum(active_websockets.values()) >= MAX_WEBSOCKETS
-            or active_websockets[ip] >= MAX_WEBSOCKETS_PER_IP):
-        await websocket.close(code=1013)  # try again later
-        return
-
-    active_websockets[ip] += 1
-    try:
-        await websocket.accept()
-        last = 0.0
-
-        while True:
-            message = await websocket.receive()
-            if message["type"] == "websocket.disconnect":
-                break
-
-            data = message.get("bytes")
-            if data is None or len(data) > MAX_FRAME_BYTES:
-                await websocket.close(code=1009 if data else 1003)
-                break
-
-            # Drop frames that arrive too fast or while the server is busy
-            # instead of queueing them.
-            now = time.monotonic()
-            if now - last < FACE_MIN_INTERVAL:
-                continue
-            last = now
-
-            try:
-                boxes = await gate.run(detect_faces, data)
-            except Busy:
-                continue
-            except InvalidImage:
-                await websocket.close(code=1003)
-                break
-
-            await websocket.send_json({"bboxes": boxes})
-
-    except WebSocketDisconnect:
-        pass
-    except Exception:
-        logger.exception("predict_face failed")
-        try:
-            await websocket.close(code=1011)
-        except Exception:
-            pass
-    finally:
-        active_websockets[ip] -= 1
-        if active_websockets[ip] <= 0:
-            del active_websockets[ip]
-
-
-# --------------------------------------------------------------------------
-# Reviews
-# --------------------------------------------------------------------------
-
-def get_api_key(authorization: str = Header(None)):
-    if not API_KEY or not authorization or not secrets.compare_digest(
-        authorization.encode(), f"Bearer {API_KEY}".encode()
-    ):
-        raise HTTPException(status_code=401, detail="Unauthorized")
-    return authorization
-
-
-class Review(BaseModel):
-    website: str = Field(min_length=1, max_length=100)
-    rating: int = Field(ge=1, le=5)
-    text: str = Field(max_length=1000)
-    # The client also sends a "date"; it is ignored, the server sets the time.
-
-
-def connect_db() -> sqlite3.Connection:
-    conn = sqlite3.connect(DATABASE, timeout=5)
-    conn.row_factory = sqlite3.Row
-    return conn
-
-
-@app.get("/review")
-def get_reviews(
-    limit: int = Query(100, ge=1, le=500),
-    offset: int = Query(0, ge=0),
-    auth=Depends(get_api_key),
-):
-    with closing(connect_db()) as conn:
-        rows = conn.execute(
-            "SELECT * FROM reviews ORDER BY id DESC LIMIT ? OFFSET ?", (limit, offset)
-        ).fetchall()
-    return {"data": [dict(row) for row in rows]}
-
-
-# Public on purpose: no credentials needed, only rate limited per IP.
-@app.post("/review", dependencies=[Depends(rate_limit(review_limiter))])
-def post_review(review: Review):
-    created_at = datetime.now(timezone.utc).isoformat()
-    with closing(connect_db()) as conn:
-        conn.execute(
-            "INSERT INTO reviews (website, rating, text, created_at) VALUES (?, ?, ?, ?)",
-            (review.website, review.rating, review.text.strip(), created_at),
-        )
-        conn.commit()
-    return {"status": "ok"}
-
-
-@app.delete("/review")
-def delete_review(auth=Depends(get_api_key)):
-    with closing(connect_db()) as conn:
-        conn.execute("DELETE FROM reviews")
-        conn.commit()
-    return {"status": "ok"}
-
-
-@app.get("/health")
-def health():
-    return {"status": "ok"}

api/src/production/settings.py

@@ -0,0 +1,57 @@
+"""All server settings in one place. Values marked (env) can be set with an
+environment variable of the same name, e.g. in compose.yaml."""
+
+import os
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+API_ROOT = Path(__file__).resolve().parents[2]    # api/ locally, /app in Docker
+
+BUILD_PATH = Path(os.getenv("BUILD_PATH", API_ROOT / "build_models"))     # (env)
+DATABASE_DIR = Path(os.getenv("DATABASE_DIR", API_ROOT / "database"))    # (env)
+DATABASE = DATABASE_DIR / "reviews.db"
+
+load_dotenv(DATABASE_DIR / ".env")
+API_KEY = os.getenv("API_KEY")  # (env) if missing, protected endpoints reject every request
+
+ORIGINS = [
+    "https://marvinkrausser.com",
+]
+
+# ---- Models ----
+BIRD_MODEL_FILE = "bird_cnn"
+BIRD_IMAGE_SIZE = 64
+BIRD_C_HIDDEN = 16
+
+FACE_MODEL_FILE = "face_detection_yolo"
+FACE_IMAGE_SIZE = 64
+FACE_GRID = 6
+FACE_C_HIDDEN = 16
+FACE_SIZE_ACTIVATION = "exp"   # must match how the deployed weights were trained
+FACE_THRESHOLD = 0.9
+
+# ---- Resource limits (the server is weak, keep everything small and bounded) ----
+# Number of predictions (bird or face) that may run at the same time. (env)
+PARALLEL_INFERENCES = max(1, int(os.getenv("PARALLEL_INFERENCES", "1")))
+# Extra requests allowed to wait for a free slot; anything beyond is rejected. (env)
+QUEUED_INFERENCES = max(0, int(os.getenv("QUEUED_INFERENCES", "1")))
+# PyTorch threads used by EACH running prediction. Total CPU use is roughly
+# PARALLEL_INFERENCES * TORCH_THREADS, so keep the product <= your CPU cores. (env)
+TORCH_THREADS = max(1, int(os.getenv("TORCH_THREADS", "1")))
+
+MAX_UPLOAD_BYTES = int(max(1, float(os.getenv("MAX_UPLOAD_MEGABYTES", "15"))) * 1024 * 1024)  # (env)
+MAX_FRAME_BYTES = 1 * 1024 * 1024
+MAX_IMAGE_PIXELS = 20_000_000
+DECODE_DRAFT_SIZE = (256, 256)  # JPEG decodes at reduced scale, still larger than the 64px model input
+
+MAX_WEBSOCKETS = 3
+MAX_WEBSOCKETS_PER_IP = 1
+FACE_MAX_FPS = float(os.getenv("FACE_MAX_FPS", "5.5"))  # (env) per websocket connection
+FACE_MIN_INTERVAL = 1 / FACE_MAX_FPS
+
+# ---- Rate limits: at most LIMIT requests per WINDOW seconds per IP ----
+PREDICT_RATE_LIMIT = max(1, int(os.getenv("PREDICT_RATE_LIMIT", "10")))          # (env)
+PREDICT_RATE_WINDOW = max(1.0, float(os.getenv("PREDICT_RATE_WINDOW", "60")))    # (env)
+REVIEW_RATE_LIMIT = max(1, int(os.getenv("REVIEW_RATE_LIMIT", "5")))             # (env)
+REVIEW_RATE_WINDOW = max(1.0, float(os.getenv("REVIEW_RATE_WINDOW", "60")))      # (env)

api/src/production/start_server.bat

@@ -1 +0,0 @@
-python -m uvicorn server:app --ws-max-size 1048576 --reload

api/src/production/start_server.sh

@@ -1 +0,0 @@
-uvicorn server:app --host 0.0.0.0 --port 8000 --ws-max-size 1048576

api/src/production/tasks.py

@@ -0,0 +1,16 @@
+"""API server tasks. Run with: python -m src.main server.<task>"""
+
+from src.production.config import Config
+
+
+def run(cfg: Config):
+    """Start the API server with uvicorn."""
+    import uvicorn
+
+    uvicorn.run("src.production.app:app", host=cfg.host, port=cfg.port,
+                reload=cfg.reload, ws_max_size=cfg.ws_max_size)
+
+
+TASKS = {
+    "run": run,
+}

api/src/production/yolo_model_production.py

@@ -1,163 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-def convert_prediction(label, image, threshold=0.9):
-    image = image.clone().detach()
-    label = label.clone().detach()
-
-    image_size = image.shape[1]
-    grid_number = label.shape[0]
-    grid_size = image_size / grid_number
-
-    boxes_to_draw = []
-    grids_to_draw_obj = []
-    grids_to_draw_noobj = []
-    for x in range(label.shape[0]):
-        for y in range(label.shape[1]):
-            if label[x, y, 4].item() < threshold:
-                grids_to_draw_noobj.append([x*grid_size, y*grid_size, (x+1)*grid_size, (y+1)*grid_size]) #xmin, ymin, xmax, ymax
-                continue
-            grids_to_draw_obj.append([x*grid_size, y*grid_size, (x+1)*grid_size, (y+1)*grid_size]) #xmin, ymin, xmax, ymax
-
-            boxx = label[x, y, 0] * (image_size / grid_number)
-            boxy = label[x, y, 1] * (image_size / grid_number)
-
-            boxw = label[x, y, 2] * image_size
-            boxh = label[x, y, 3] * image_size
-
-            boxx, boxy = turn_image_centered(x=boxx, y=boxy, img_w=image_size, img_h=image_size, S=grid_number, cell_i=x, cell_j=y)
-
-            boxes_to_draw.append([label[x, y, 4]] + xy_center_to_edges(boxx, boxy, boxw, boxh)) #xmin, ymin, xmax, ymax
-    boxes_to_draw = nms(boxes_to_draw)
-    return boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj
-
-def turn_image_centered(x, y, img_w, img_h, S, cell_i, cell_j):
-    cell_w = img_w / S
-    cell_h = img_h / S
-
-    cell_border_w = cell_w * cell_i
-    cell_border_h = cell_h * cell_j
-
-    return x + cell_border_w, y + cell_border_h
-
-def xy_center_to_edges(xcenter, ycenter, width, height):
-    width = max(width, 1)
-    height = max(height, 1)
-
-    x = xcenter - (width / 2)
-    y = ycenter - (height / 2)
-
-    return [x, y, x + width, y + height]
-
-def iou(boxA, boxB):
-    xA = max(boxA[0], boxB[0])
-    yA = max(boxA[1], boxB[1])
-    xB = min(boxA[2], boxB[2])
-    yB = min(boxA[3], boxB[3])
-
-    inter_area = max(0, xB - xA) * max(0, yB - yA)
-
-    boxA_area = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
-    boxB_area = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
-
-    union = boxA_area + boxB_area - inter_area
-
-    return inter_area / union if union > 0 else 0
-
-def nms(bboxes, iou_threshold=0.1):
-    bboxes = sorted(bboxes, key=lambda x: x[0], reverse=True)
-
-    keep = []
-
-    while bboxes:
-        best = bboxes.pop(0)
-        keep.append(best[1:5])
-
-        bboxes = [
-            box for box in bboxes
-            if iou(best[1:5], box[1:5]) < iou_threshold
-        ]
-
-    return keep
-
-class Yolo_Conv_Block(nn.Module):
-    def __init__(self, c_in, c_hidden, c_out, kernel_size):
-        super().__init__()
-
-        self.model = nn.Sequential(
-            nn.Conv2d(in_channels=c_in, out_channels=c_hidden, kernel_size=kernel_size, padding=kernel_size//2),
-            nn.BatchNorm2d(c_hidden),
-            nn.LeakyReLU(),
-            nn.Conv2d(in_channels=c_hidden, out_channels=c_out, kernel_size=1)
-        )
-
-    def forward(self, x):
-        return self.model(x)
-    
-class SkipBlock(nn.Module):
-    def __init__(self, c_in, c_out, kernel_size=3):
-        super().__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size//2),
-            nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
-            nn.LeakyReLU(inplace=True),
-
-            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
-            nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
-            nn.LeakyReLU(inplace=True),
-
-            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
-            nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
-            nn.LeakyReLU(inplace=True)
-        )
-        self.conv_skip = nn.Sequential(
-            nn.Conv2d(c_in, c_out, 1),
-            nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
-            nn.LeakyReLU(inplace=True)
-        )
-
-    def forward(self, x):
-        return(F.dropout(F.leaky_relu(self.conv_skip(x) + self.conv(x), inplace=True), p=0.3))
-
-
-class Yolo_model(nn.Module):
-    def __init__(self, c_in, boxes, grid, labels, c_hidden=16):
-        super().__init__()
-
-        self.model = nn.Sequential(
-            nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups=c_hidden//8, num_channels=c_hidden),
-            nn.LeakyReLU(inplace=True),
-
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            
-            SkipBlock(c_in=c_hidden, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*4),
-            SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
-            SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
-            SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
-
-            nn.Conv2d(c_hidden*4, c_hidden*8, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups=c_hidden//2, num_channels=c_hidden*8),
-            nn.LeakyReLU(inplace=True),
-            nn.Dropout(0.3),
-
-            nn.AdaptiveAvgPool2d((grid, grid)),
-            nn.Conv2d(c_hidden*8, boxes*5 + labels, kernel_size=1)
-        )
-
-    def forward(self, x):
-        x = self.model(x).permute(0, 2, 3, 1)
-        center = F.sigmoid(x[..., :2])
-        size = torch.exp(x[..., 2:4])
-        conf_class = F.sigmoid(x[..., 4:])
-        return torch.cat([center, size, conf_class], dim=3)

api/src/r_cnn/config.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Config:
+    """Settings for the R-CNN style experiments (classify crops, then run
+    selective search on a full image). Paths are relative to api/.
+    Every field can be overridden on the command line, e.g. --epochs 20."""
+
+    # Datasets
+    faces_raw_dir: str = "data/faces/train"             # COCO: source for create-dataset
+    faces_processed_dir: str = "data/faces_processed"   # ImageFolder: face/ and background/
+    football_train_dir: str = "data/football/train"     # COCO
+    football_val_dir: str = "data/football/valid"       # COCO
+    annotation_file: str = "_annotations.coco.json"
+
+    # Crops
+    crop_size: int = 64
+    padding: int = 40                     # football: extra pixels around each box/proposal
+    backgrounds_per_image: int = 10       # create-dataset: background crops per image
+    train_fraction: float = 0.8
+
+    # Model
+    c_hidden: int = 32
+    layers: int = 10
+    num_classes: int = 2
+
+    # Training
+    epochs: int = 200
+    batch_size: int = 256                 # faces (pre-cropped images)
+    football_batch_size: int = 2          # football (full images, cropped on the fly)
+    learning_rate: float = 1e-3
+    weight_decay: float = 1e-4
+    background_weight: float = 1 / 3      # class weight of "background" in the faces loss
+    device: str = "auto"
+
+    # Output
+    save_dir: str = "saved_models"
+    faces_model_name: str = "face_detection"
+    football_model_name: str = "object_detection"
+    save: bool = True
+
+    # detect tasks: proposals smaller than min_size or larger than max_size are ignored
+    image: str = ""                       # detect-faces: image to run on
+    min_size: int = 5
+    max_size: int = 600
+    min_confidence: float = 0.97
+    football_max_size: int = 100
+    football_min_confidence: float = 0.8

api/src/r_cnn/createObjectDetectionDataset.py

@@ -1,120 +0,0 @@
-import random
-from pathlib import Path
-
-import torch
-import torchvision.transforms.functional as TF
-from torchvision.transforms import ToPILImage, ToTensor
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-
-from src.r_cnn.cocoDetectionDataset import CocoDetectionDataset
-from src.util import iou
-
-
-
-SAVE_PATH = "./saved_models"
-
-def get_transform():
-    return ToTensor()
-
-dataset = CocoDetectionDataset(
-    image_dir="data/faces/train", 
-    annotation_path="data/faces/train/_annotations.coco.json",
-    transforms=get_transform()
-)
-
-processing_loader = DataLoader(dataset, batch_size=1, shuffle=False, collate_fn=lambda x: tuple(zip(*x)))
-
-def random_crop(W, H, sizeX=64, sizeY =64):
-    x = random.randint(0, max(W - sizeX, 0))
-    y = random.randint(0, max(H - sizeY, 0))
-    return x, y, x + sizeX, y + sizeY
-
-def box_inside(inner_box, outer_box):
-    return (
-        inner_box[0] >= outer_box[0] and
-        inner_box[1] >= outer_box[1] and
-        inner_box[2] <= outer_box[2] and
-        inner_box[3] <= outer_box[3]
-    )
-
-def is_background(crop, gt_boxes, threshold=0.0):
-    for box in gt_boxes:
-        if iou(crop, box) > threshold:
-            return False
-    return True
-
-def get_background_crop(image, gt_boxes, max_trials=100):
-    C, H, W = image.shape
-
-    for _ in range(max_trials):
-        x1, y1, x2, y2 = random_crop(W, H, sizeX=random.randint(50, 600), sizeY=random.randint(50, 600))
-
-        crop_box = (x1, y1, x2, y2)
-
-        if is_background(crop_box, gt_boxes, threshold=0.0):
-            crop = image[:, y1:y2, x1:x2]
-            return crop
-
-    return None
-
-def create_background_tensor(amount, dataset, labels, boxes, image):
-    for _ in range(amount):
-        background = get_background_crop(image=image, gt_boxes=boxes.to(torch.int64))
-        if background is None:
-            break
-        background = TF.resize(background, [64, 64], antialias=True)
-        dataset.append(background)
-        labels.append(torch.tensor(0))
-    return dataset, labels
-
-def create_stack(images, annotations, PADDING):
-    croped_images = []
-    labels = []
-    for i in range(len(images)):
-        image = images[i]
-        annotation = annotations[i]
-        _, h, w = image.shape
-        
-        for box, label in tuple(zip(annotation["boxes"], annotation["labels"])):
-            box = box.to(torch.int64)
-            box_copy = []
-
-            box_copy.append(max(box[0]-PADDING, 0))
-            box_copy.append(max(box[1]-PADDING, 0))
-            box_copy.append(min(box[2]+PADDING, w-1))
-            box_copy.append(min(box[3]+PADDING, h-1))
-
-            if box_copy[1] == box_copy[3] or box_copy[0] == box_copy[2]:
-                continue
-
-            croped_image = image[: , box_copy[1]:box_copy[3], box_copy[0]:box_copy[2]]
-            croped_image = TF.resize(croped_image, [64, 64], antialias=True)
-            croped_images.append(croped_image)
-            labels.append(label)
-
-        croped_images, labels = create_background_tensor(amount=10, dataset=croped_images, labels=labels, boxes=annotation["boxes"], image=image)
-    return croped_images, labels
-
-def create_croped_dataset():
-    to_pil = ToPILImage()
-
-    out_path = Path("data/faces_processed/face")
-    out_path.mkdir(parents=True, exist_ok=True)
-    out_path = Path("data/faces_processed/background")
-    out_path.mkdir(parents=True, exist_ok=True)
-
-    count_image = 0
-    count_background = 0
-    for images, annotations in tqdm(processing_loader):
-        croped_images, labels = create_stack(images, annotations, PADDING=0)
-
-        for i in range(len(croped_images)):
-            pil_img = to_pil(croped_images[i])
-
-            if labels[i] == 1:
-                pil_img.save(f"data/faces_processed/face/{count_image}.jpg")
-                count_image += 1
-            else:
-                pil_img.save(f"data/faces_processed/background/{count_background}.jpg")
-                count_background += 1

api/src/r_cnn/crops.py

@@ -0,0 +1,59 @@
+"""Cutting object and background crops out of detection images."""
+
+import random
+
+import torch
+import torchvision.transforms.functional as TF
+
+from src.common.boxes import iou
+
+
+def random_crop(W, H, sizeX=64, sizeY=64):
+    x = random.randint(0, max(W - sizeX, 0))
+    y = random.randint(0, max(H - sizeY, 0))
+    return x, y, x + sizeX, y + sizeY
+
+
+def is_background(crop, gt_boxes, threshold=0.0):
+    return all(iou(crop, box) <= threshold for box in gt_boxes)
+
+
+def get_background_crop(image, gt_boxes, max_trials=100, iou_threshold=0.0):
+    """A random crop that overlaps no ground-truth box, or None."""
+    _, H, W = image.shape
+    for _ in range(max_trials):
+        x1, y1, x2, y2 = random_crop(W, H, sizeX=random.randint(50, 600), sizeY=random.randint(50, 600))
+        if is_background((x1, y1, x2, y2), gt_boxes, threshold=iou_threshold):
+            return image[:, y1:y2, x1:x2]
+    return None
+
+
+def create_stack(images, annotations, padding, crop_size=64, backgrounds_per_image=0):
+    """Crops every annotated box (label from the annotation) plus
+    `backgrounds_per_image` background crops (label 0) per image."""
+    crops, labels = [], []
+
+    for image, annotation in zip(images, annotations):
+        _, h, w = image.shape
+
+        for box, label in zip(annotation["boxes"], annotation["labels"]):
+            box = box.to(torch.int64)
+            x1 = max(box[0] - padding, 0)
+            y1 = max(box[1] - padding, 0)
+            x2 = min(box[2] + padding, w - 1)
+            y2 = min(box[3] + padding, h - 1)
+
+            if x1 == x2 or y1 == y2:   # zero-size box, cannot be resized
+                continue
+
+            crops.append(TF.resize(image[:, y1:y2, x1:x2], [crop_size, crop_size], antialias=True))
+            labels.append(label)
+
+        for _ in range(backgrounds_per_image):
+            background = get_background_crop(image, annotation["boxes"].to(torch.int64))
+            if background is None:
+                break
+            crops.append(TF.resize(background, [crop_size, crop_size], antialias=True))
+            labels.append(torch.tensor(0))
+
+    return crops, labels

api/src/r_cnn/detect.py

@@ -0,0 +1,58 @@
+"""Selective search + crop classifier on a full image."""
+
+import cv2
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from tqdm import tqdm
+
+
+def resize_keep_aspect(img, target_w, target_h):
+    h, w = img.shape[:2]
+    scale = min(target_w / w, target_h / h)
+    new_w, new_h = int(w * scale), int(h * scale)
+    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
+    return resized, w / new_w, h / new_h
+
+
+def detect(model, image, device, padding=0, crop_size=64, min_size=5, max_size=600, min_confidence=0.8):
+    """image: tensor [C, H, W] in 0..1. Shows the image, then the detections.
+    Needs opencv-contrib-python (cv2.ximgproc)."""
+    if not hasattr(cv2, "ximgproc"):
+        raise RuntimeError("Selective search needs opencv-contrib-python (cv2.ximgproc)")
+
+    _, H, W = image.shape
+    image_numpy = (image.permute(1, 2, 0).cpu().numpy() * 255).astype("uint8")
+    image_numpy = cv2.cvtColor(image_numpy, cv2.COLOR_BGR2RGB)
+    resized, scale_w, scale_h = resize_keep_aspect(image_numpy, W, H)
+
+    cv2.imshow("", resized)
+    cv2.waitKey(0)
+
+    ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()
+    ss.setBaseImage(resized)
+    ss.switchToSelectiveSearchFast()
+
+    predictions = []
+    with torch.inference_mode():
+        for (x, y, w, h) in tqdm(ss.process()):
+            x, y, w, h = int(x * scale_w), int(y * scale_h), int(w * scale_w), int(h * scale_h)
+            if w < min_size or h < min_size or w > max_size or h > max_size:
+                continue
+
+            x1, y1 = max(x - padding, 0), max(y - padding, 0)
+            x2, y2 = min(x + w + padding, W - 1), min(y + h + padding, H - 1)
+
+            crop = TF.resize(image[:, y1:y2, x1:x2], [crop_size, crop_size], antialias=True)
+            probs = F.softmax(model(crop.unsqueeze(0).to(device)), dim=1)
+            confidence, cls = torch.max(probs, dim=1)
+            if cls.item() == 1 and confidence.item() > min_confidence:
+                predictions.append((confidence.item(), (x, y, w, h)))
+
+    print(f"{len(predictions)} detections")
+    draw = image_numpy.copy()
+    for _, (x, y, w, h) in sorted(predictions, reverse=True):
+        cv2.rectangle(draw, (x, y), (x + w, y + h), (255, 0, 0), 1)
+
+    cv2.imshow("", draw)
+    cv2.waitKey(0)

api/src/r_cnn/r_cnn.py

@@ -1,315 +0,0 @@
-import os
-import random
-
-import cv2
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchvision.transforms.functional as TF
-from tqdm import tqdm
-
-from src.util import iou, visualizeImage
-
-
-class ObjectDetectionCNN(nn.Module):
-    def __init__(self, c_in, c_hidden, c_out, layers):
-        super().__init__()
-
-        self.model = nn.ModuleList()
-
-        self.model.append(nn.Sequential(
-                nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
-                nn.BatchNorm2d(c_hidden),
-                nn.ReLU(inplace=True)
-            ))
-
-        for _ in range(layers-1):
-            self.model.append(nn.Sequential(
-                nn.Conv2d(c_hidden, c_hidden, kernel_size=3, padding=1),
-                nn.BatchNorm2d(c_hidden),
-                nn.ReLU(inplace=True)
-            ))
-
-        self.model.append(nn.Sequential(
-            nn.AdaptiveAvgPool2d((1, 1)),
-            nn.Flatten(),
-            nn.Linear(c_hidden, c_out),
-            nn.Dropout(0.3)
-        ))
-
-    def forward(self, x):
-        for layer in self.model:
-            x = layer(x)
-        return x
-    
-
-
-def random_crop(W, H, sizeX=64, sizeY =64):
-    x = random.randint(0, max(W - sizeX, 0))
-    y = random.randint(0, max(H - sizeY, 0))
-    return x, y, x + sizeX, y + sizeY
-
-def is_background(crop, gt_boxes, threshold=0.0):
-    for box in gt_boxes:
-        if iou(crop, box) > threshold:
-            return False
-    return True
-
-def get_background_crop(image, gt_boxes, max_trials=100):
-    C, H, W = image.shape
-
-    for _ in range(max_trials):
-        x1, y1, x2, y2 = random_crop(W, H, sizeX=random.randint(50, 600), sizeY=random.randint(50, 600))
-
-        crop_box = (x1, y1, x2, y2)
-
-        if is_background(crop_box, gt_boxes, threshold=0.1):
-            crop = image[:, y1:y2, x1:x2]
-            return crop
-
-    return None
-
-def create_background_tensor(amount, dataset, labels, boxes, image):
-    for _ in range(amount):
-                    background = get_background_crop(image=image, gt_boxes=boxes.to(torch.int64))
-                    if background is not None:
-                        background = TF.resize(background, [64, 64], antialias=True)
-                        dataset.append(background)
-                        labels.append(torch.tensor(0))
-                    else:
-                         print("Background not found")
-    return dataset, labels
-
-def create_stack(images, annotations, PADDING):
-    croped_images = []
-    labels = []
-    for i in range(len(images)):
-        image = images[i]
-        annotation = annotations[i]
-        _, h, w = image.shape
-        
-        for box, label in tuple(zip(annotation["boxes"], annotation["labels"])):
-            box = box.to(torch.int64)
-            box_copy = []
-
-            box_copy.append(max(box[0]-PADDING, 0))
-            box_copy.append(max(box[1]-PADDING, 0))
-            box_copy.append(min(box[2]+PADDING, w-1))
-            box_copy.append(min(box[3]+PADDING, h-1))
-            
-
-            croped_image = image[: , box_copy[1]:box_copy[3], box_copy[0]:box_copy[2]]
-            #visualizeImage(croped_image)
-            croped_image = TF.resize(croped_image, [64, 64], antialias=True)
-            #visualizeImage(croped_image)
-            croped_images.append(croped_image)
-            labels.append(label)
-
-        #croped_images, labels = create_background_tensor(amount=len(croped_images), dataset=croped_images, labels=labels, boxes=annotation["boxes"], image=image)
-    return croped_images, labels
-
-def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True, PADDING=20, device="cpu"):
-    best_val = torch.finfo(torch.float32).max
-
-    for epoch in range(200):
-        ############
-        # Training #
-        ############
-        model.train()
-
-        true_preds, count, lossCount = 0, 0, 0.
-        for images, annotations in tqdm(train_loader, desc=f"Train", leave=False):
-            croped_images, labels = create_stack(images, annotations, PADDING)
-
-            croped_images = torch.stack(croped_images).to(device)
-            labels = torch.stack(labels).to(device)
-
-            prediction = model(croped_images)
-
-            loss = loss_module(prediction, labels)
-            lossCount += loss.sum().item()
-
-            optimizer.zero_grad()
-
-            loss.backward()
-
-            optimizer.step()
-
-            true_preds += (prediction.argmax(dim=1) == labels).sum().item()
-            count += croped_images.size(0)
-
-        train_acc = true_preds / count
-        train_loss = lossCount / count
-
-        torch.cuda.empty_cache()
-
-        ##############
-        # Validation #
-        ##############
-        model.eval()
-
-        true_preds, count, lossCount = 0, 0, 0.
-        for images, annotations in tqdm(val_loader, desc=f"Test", leave=False):
-            with torch.no_grad():
-                croped_images, labels = create_stack(images, annotations, PADDING)
-
-                croped_images = torch.stack(croped_images).to(device)
-                labels = torch.stack(labels).to(device)
-
-                prediction = model(croped_images)
-
-                loss = loss_module(prediction, labels)
-                lossCount += loss.sum().item()
-
-                true_preds += (prediction.argmax(dim=1) == labels).sum().item()
-                count += croped_images.size(0)
-
-        val_acc = true_preds / count
-        val_loss = lossCount / count
-
-        if(saving and best_val > val_loss):
-                best_val = val_loss
-                save_dir = os.path.join(SAVE_PATH, model_name)
-                os.makedirs(save_dir, exist_ok=True)
-
-                save_path = os.path.join(save_dir, model_name)
-                torch.save(model.state_dict(), save_path)
-
-        print(f"epoch: {epoch+1} | train accuracy: {int(train_acc * 1000) / 10}% | validation accuracy: {int(val_acc * 1000) / 10}% | train loss: {int(train_loss * 1000) / 100} | val loss: {int(val_loss * 1000) / 100}")
-        torch.cuda.empty_cache()
-    return best_val
-
-
-def trainNormalDataset(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True, device="cpu"):
-    best_val = torch.finfo(torch.float32).max
-
-    for epoch in range(200):
-        ############
-        # Training #
-        ############
-        model.train()
-
-        true_preds, count, lossCount = 0, 0, 0.
-        for images, labels in tqdm(train_loader, desc=f"Train", leave=False):
-            images = images.to(device)
-            labels = labels.to(device)
-
-            prediction = model(images)
-
-            loss = loss_module(prediction, labels)
-            lossCount += loss.sum().item()
-
-            optimizer.zero_grad()
-
-            loss.backward()
-
-            optimizer.step()
-
-            true_preds += (prediction.argmax(dim=1) == labels).sum().item()
-            count += images.size(0)
-
-        train_acc = true_preds / count
-        train_loss = lossCount / count
-
-        torch.cuda.empty_cache()
-
-        ##############
-        # Validation #
-        ##############
-        model.eval()
-
-        true_preds, count, lossCount = 0, 0, 0.
-        for images, labels in tqdm(val_loader, desc=f"Test", leave=False):
-            with torch.no_grad():
-                images = images.to(device)
-                labels = labels.to(device)
-
-                prediction = model(images)
-
-                loss = loss_module(prediction, labels)
-                lossCount += loss.sum().item()
-
-                true_preds += (prediction.argmax(dim=1) == labels).sum().item()
-                count += images.size(0)
-
-        val_acc = true_preds / count
-        val_loss = lossCount / count
-
-        if(saving and best_val > val_loss):
-                best_val = val_loss
-                save_dir = os.path.join(SAVE_PATH, model_name)
-                os.makedirs(save_dir, exist_ok=True)
-
-                save_path = os.path.join(save_dir, model_name)
-                torch.save(model.state_dict(), save_path)
-
-        print(f"epoch: {epoch+1} | train accuracy: {int(train_acc * 1000) / 10}% | validation accuracy: {int(val_acc * 1000) / 10}% | train loss: {int(train_loss * 1000) / 100} | val loss: {int(val_loss * 1000) / 100}")
-        torch.cuda.empty_cache()
-    return best_val
-
-def resize_keep_aspect(img, target_w, target_h):
-    h, w = img.shape[:2]
-
-    scale = min(target_w / w, target_h / h)
-    new_w = int(w * scale)
-    new_h = int(h * scale)
-
-    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
-    return resized, w / new_w, h / new_h
-
-def map_box_to_original(box, scale_x, scale_y):
-    x1, y1, x2, y2 = box
-
-    return int(x1 * scale_x), int(y1 * scale_y), int(x2 * scale_x), int(y2 * scale_y)
-
-def eval(model, image, BUILD_PATH, device, PADDING = 0, minSize=5, maxSize=600, minConf=0.8):
-    _, H, W = image.shape
-
-    model.load_state_dict(torch.load(BUILD_PATH, map_location=torch.device(device)))
-    model.to(device)
-    model.eval()
-    image_numpy = image.permute(1, 2, 0).cpu().numpy()
-    image_numpy = (image_numpy*255).astype("uint8")
-    image_numpy = cv2.cvtColor(image_numpy, cv2.COLOR_BGR2RGB)
-    image_numpy_copy, resizedW, resizedH = resize_keep_aspect(image_numpy, W, H)
-
-    cv2.imshow("", image_numpy_copy)
-    cv2.waitKey(0)
-
-    ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()
-    ss.setBaseImage(image_numpy_copy)
-    ss.switchToSelectiveSearchFast()
-
-    rects = ss.process()
-
-    draw = image_numpy.copy()
-    predictions = []
-
-    for (x, y, w, h) in tqdm(rects):
-        x, y, w, h = map_box_to_original((x, y, w, h), resizedW, resizedH)
-
-        if w < minSize or h < minSize or w > maxSize or h > maxSize:
-            continue
-
-        x1 = max(x-PADDING, 0)
-        y1 = max(y-PADDING, 0)
-        x2 = min(x+w+PADDING, W-1)
-        y2 = min(y+h+PADDING, H-1)
-
-        crop = image[:, y1:y2, x1:x2]
-        crop = TF.resize(crop, [64, 64], antialias=True)
-        crop = crop.unsqueeze(0).to(device)
-
-        pred = model(crop)
-        probs = F.softmax(pred, dim=1)
-        confidence, cls = torch.max(probs, dim=1)
-        if cls.item() == 1 and confidence.item() > minConf:
-            predictions.append((confidence.item(), cls.item(), (x, y, w, h)))
-
-    predictions = sorted(predictions, key=lambda x: x[0], reverse=True)
-    print(len(predictions))
-    for conf, cls, (x, y, w, h) in predictions[:]:
-        cv2.rectangle(draw, (x, y), (x + w, y + h), (255, 0, 0), 1)
-
-    cv2.imshow("", draw)
-    cv2.waitKey(0)

api/src/r_cnn/r_cnn_test.py

@@ -1,50 +0,0 @@
-import os
-from torch import tensor
-import torch
-from torch.utils.data import DataLoader
-from torchvision.transforms import ToTensor
-import torch.nn as nn
-
-from src.r_cnn.cocoDetectionDataset import CocoDetectionDataset
-from src.r_cnn.r_cnn import ObjectDetectionCNN, eval, train
-
-
-def get_transform():
-    return ToTensor()
-
-def train_cnn_test():
-
-    SAVE_PATH = "./saved_models"
-    PADDING = 20
-
-    train_dataset = CocoDetectionDataset(
-        image_dir="data/football/train", 
-        annotation_path="data/football/train/_annotations.coco.json",
-        transforms=get_transform()
-    )
-
-    val_dataset = CocoDetectionDataset(
-        image_dir="data/football/valid",
-        annotation_path="data/football/valid/_annotations.coco.json",
-        transforms=get_transform()
-    )
-
-    train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
-    val_loader = DataLoader(val_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
-
-    device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
-    print("Using device", device)
-
-    model = ObjectDetectionCNN(c_in=3, c_hidden=32, c_out=2, layers=10)
-    model.to(device)
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
-    loss_module = nn.CrossEntropyLoss()
-
-    #train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, 
-    #                optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, PADDING=40, device=device)
-
-    #exit()
-
-    image = next(iter(val_loader))[0][0]
-    eval(model=model, image=image, BUILD_PATH=os.path.join(SAVE_PATH, "object_detection", "object_detection"), 
-        device=device, PADDING=40, minSize=5, maxSize=100, minConf=0.8)

api/src/r_cnn/tasks.py

@@ -0,0 +1,126 @@
+"""R-CNN style experiments. Run with: python -m src.main rcnn.<task>"""
+
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from torchvision.io import read_image
+from torchvision.transforms import ToPILImage, ToTensor
+from tqdm import tqdm
+
+from src.common.checkpoints import checkpoint_path, load_checkpoint
+from src.common.data import collate_as_tuples, split_dataset
+from src.common.device import get_device
+from src.common.training import fit
+from src.models.r_cnn import ObjectDetectionCNN
+from src.r_cnn.config import Config
+from src.r_cnn.crops import create_stack
+from src.r_cnn.dataset import CocoDetectionDataset
+from src.r_cnn.detect import detect
+
+
+def build_model(cfg: Config):
+    return ObjectDetectionCNN(c_in=3, c_hidden=cfg.c_hidden, c_out=cfg.num_classes, layers=cfg.layers)
+
+
+def coco_dataset(cfg: Config, folder: str):
+    return CocoDetectionDataset(image_dir=folder,
+                                annotation_path=os.path.join(folder, cfg.annotation_file),
+                                transforms=ToTensor())
+
+
+def create_dataset(cfg: Config):
+    """Cut face and background crops out of --faces-raw-dir (for train-faces)."""
+    loader = DataLoader(coco_dataset(cfg, cfg.faces_raw_dir), batch_size=1, shuffle=False,
+                        collate_fn=collate_as_tuples)
+    to_pil = ToPILImage()
+    out = Path(cfg.faces_processed_dir)
+    (out / "face").mkdir(parents=True, exist_ok=True)
+    (out / "background").mkdir(parents=True, exist_ok=True)
+
+    counts = {"face": 0, "background": 0}
+    for images, annotations in tqdm(loader):
+        crops, labels = create_stack(images, annotations, padding=0, crop_size=cfg.crop_size,
+                                     backgrounds_per_image=cfg.backgrounds_per_image)
+        for crop, label in zip(crops, labels):
+            kind = "face" if label == 1 else "background"
+            to_pil(crop).save(out / kind / f"{counts[kind]}.jpg")
+            counts[kind] += 1
+
+    print(f"Saved {counts['face']} faces and {counts['background']} backgrounds to {out}")
+
+
+def train_faces(cfg: Config):
+    """Train the face/background crop classifier on --faces-processed-dir."""
+    device = get_device(cfg.device)
+
+    base = [transforms.Resize(cfg.crop_size), transforms.CenterCrop(cfg.crop_size), transforms.ToTensor()]
+    augment = [transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(degrees=15)] + base
+    train_ds, val_ds = split_dataset(datasets.ImageFolder(cfg.faces_processed_dir), cfg.train_fraction,
+                                     transforms.Compose(augment), transforms.Compose(base))
+    train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True)
+    val_loader = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False)
+
+    model = build_model(cfg).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
+    # ImageFolder sorts classes alphabetically: 0 = background, 1 = face
+    weights = torch.tensor([cfg.background_weight, 1.0], device=device)
+
+    fit(model, nn.CrossEntropyLoss(weight=weights), optimizer, train_loader, val_loader, device,
+        epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.faces_model_name, save=cfg.save,
+        monitor="loss", track_accuracy=True)
+
+
+def detect_faces(cfg: Config):
+    """Run selective search + the face classifier on --image."""
+    if not cfg.image:
+        raise SystemExit("Pass the image to run on with --image path/to/image.jpg")
+    device = get_device(cfg.device)
+    model = load_checkpoint(build_model(cfg), checkpoint_path(cfg.save_dir, cfg.faces_model_name), device)
+    image = read_image(cfg.image).float() / 255.0
+    detect(model, image, device, crop_size=cfg.crop_size, min_size=cfg.min_size,
+           max_size=cfg.max_size, min_confidence=cfg.min_confidence)
+
+
+def train_football(cfg: Config):
+    """Train the crop classifier on the football COCO dataset (crops on the fly)."""
+    device = get_device(cfg.device)
+    train_loader = DataLoader(coco_dataset(cfg, cfg.football_train_dir), batch_size=cfg.football_batch_size,
+                              shuffle=True, collate_fn=collate_as_tuples)
+    val_loader = DataLoader(coco_dataset(cfg, cfg.football_val_dir), batch_size=cfg.football_batch_size,
+                            shuffle=True, collate_fn=collate_as_tuples)
+
+    def prepare_batch(batch, device):
+        images, annotations = batch
+        crops, labels = create_stack(images, annotations, padding=cfg.padding, crop_size=cfg.crop_size)
+        return torch.stack(crops).to(device), torch.stack(labels).to(device)
+
+    model = build_model(cfg).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
+
+    fit(model, nn.CrossEntropyLoss(), optimizer, train_loader, val_loader, device,
+        epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.football_model_name, save=cfg.save,
+        monitor="loss", track_accuracy=True, prepare_batch=prepare_batch)
+
+
+def detect_football(cfg: Config):
+    """Run selective search + the football classifier on a validation image."""
+    device = get_device(cfg.device)
+    model = load_checkpoint(build_model(cfg), checkpoint_path(cfg.save_dir, cfg.football_model_name), device)
+    loader = DataLoader(coco_dataset(cfg, cfg.football_val_dir), batch_size=1, shuffle=True,
+                        collate_fn=collate_as_tuples)
+    image = next(iter(loader))[0][0]
+    detect(model, image, device, padding=cfg.padding, crop_size=cfg.crop_size, min_size=cfg.min_size,
+           max_size=cfg.football_max_size, min_confidence=cfg.football_min_confidence)
+
+
+TASKS = {
+    "create-dataset": create_dataset,
+    "train-faces": train_faces,
+    "detect-faces": detect_faces,
+    "train-football": train_football,
+    "detect-football": detect_football,
+}

api/src/r_cnn/train_face_detection_r_cnn.py

@@ -1,69 +0,0 @@
-import os
-
-import torch
-import torch.nn as nn
-from torchvision import datasets, transforms
-from torchvision.io import read_image
-from torch.utils.data import DataLoader, WeightedRandomSampler, random_split
-from tqdm import tqdm
-
-from src.r_cnn.r_cnn import ObjectDetectionCNN, eval, trainNormalDataset
-from src.util import TransformedSubset, visualizeImage
-
-SAVE_PATH = "./saved_models"
-PADDING = 20
-IMAGE_SIZE = 64
-
-transform = transforms.Compose([
-    transforms.Resize(IMAGE_SIZE),
-    transforms.CenterCrop(IMAGE_SIZE),
-    transforms.ToTensor()
-])
-
-transform_augemnt = transforms.Compose([
-    transforms.RandomHorizontalFlip(p=0.5),
-    transforms.RandomRotation(degrees=15),
-    transforms.Resize(IMAGE_SIZE),
-    transforms.CenterCrop(IMAGE_SIZE),
-    transforms.ToTensor()
-])
-
-dataset = datasets.ImageFolder("data/faces_processed")
-
-weights = [1/3, 1]
-
-sampler = WeightedRandomSampler(
-    weights=weights,
-    num_samples=len(dataset),
-    replacement=True
-)
-
-train_size = int(0.8 * len(dataset))
-val_size = len(dataset) - train_size
-
-train_subset, val_subset = random_split(dataset, [train_size, val_size])
-train_dataset = TransformedSubset(train_subset, transform_augemnt)
-val_dataset = TransformedSubset(val_subset, transform)
-
-
-train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True)
-val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False)
-
-device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
-print("Using device", device)
-
-model = ObjectDetectionCNN(c_in=3, c_hidden=32, c_out=2, layers=10)
-model.to(device)
-optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
-loss_module = nn.CrossEntropyLoss(weight=torch.tensor([1/3, 1], device=device, dtype=torch.float32))
-
-#trainNormalDataset(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, 
-#                optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, device=device, model_name="face_detection")
-
-#exit()
-
-image = read_image("data/faces/train/_url-http_3A_2F_2Fdingyue-ws-126-net_2F2023_2F0201_2F6a17ca80j00rpd7yx00bmd000dw00gop_jpg.rf.Setza3JikTSzG04c0Fd1.jpg")
-image = image.float() / 255.0
-visualizeImage(image)
-
-eval(model=model, image=image, BUILD_PATH=os.path.join(SAVE_PATH, "face_detection", "face_detection"), device=device, minConf=0.97)

api/src/util.py

@@ -1,74 +0,0 @@
-from matplotlib import pyplot as plt
-import torch
-from tqdm import tqdm
-from torch.utils.data import DataLoader
-import time
-import multiprocessing as mp
-
-class TransformedSubset(torch.utils.data.Dataset):
-    def __init__(self, subset, transform=None):
-        self.subset = subset
-        self.transform = transform
-
-    def __getitem__(self, idx):
-        x, y = self.subset[idx]
-
-        if self.transform:
-            x = self.transform(x)
-
-        return x, y
-
-    def __len__(self):
-        return len(self.subset)
-    
-
-def visualizeData(dataset):
-    images, labels = next(iter(dataset))
-
-    for i in range(4):
-        img = images[i]
-
-        img = img.permute(1, 2, 0)
-
-        plt.figure(figsize=(3,3))
-        plt.imshow(img)
-        plt.title(f"Label: {labels[i].item()}")
-        plt.axis("off")
-
-    plt.show()
-
-def visualizeImage(image):
-    image = image.permute(1, 2, 0)
-
-    plt.figure(figsize=(3,3))
-    plt.imshow(image)
-    plt.axis("off")
-
-    plt.show()
-
-def iou(boxA, boxB):
-    xA = max(boxA[0], boxB[0])
-    yA = max(boxA[1], boxB[1])
-    xB = min(boxA[2], boxB[2])
-    yB = min(boxA[3], boxB[3])
-
-    inter_area = max(0, xB - xA) * max(0, yB - yA)
-
-    boxA_area = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
-    boxB_area = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
-
-    union = boxA_area + boxB_area - inter_area
-
-    return inter_area / union if union > 0 else 0
-
-def test_workers_speed(dataset, model):
-    device = next(model.parameters()).device
-    for num_workers in range(0, mp.cpu_count(), 2):  
-        train_loader = DataLoader(dataset,shuffle=True,num_workers=num_workers,batch_size=64,pin_memory=True)
-        start = time.time()
-        for _ in range(2):
-            for images, _  in tqdm(train_loader, leave=False):
-                images = images.to(device)
-                _ = model(images)
-        end = time.time()
-        print("Finish with:{} seconds, num_workers={}".format(int(end - start), num_workers))

api/src/yolo/config.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Config:
+    """Settings for the YOLO face detector. Paths are relative to api/.
+    Every field can be overridden on the command line, e.g. --batch-size 32."""
+
+    # Data (COCO format: images + _annotations.coco.json in each folder)
+    train_dir: str = "data/faces_scenery/train"
+    val_dir: str = "data/faces_scenery/test"
+    annotation_file: str = "_annotations.coco.json"
+    augment_train: bool = True
+    augment_val: bool = True          # the original code augmented validation too
+    image_size: int = 64
+    grid: int = 6
+
+    # Model
+    c_hidden: int = 16
+    boxes: int = 1
+    labels: int = 1
+    size_activation: str = "sigmoid"  # "sigmoid" (current training) or "exp" (deployed model)
+
+    # Training
+    epochs: int = 500
+    batch_size: int = 64
+    learning_rate: float = 1e-3
+    weight_decay: float = 1e-4
+    train_workers: int = 2
+    val_workers: int = 0
+    device: str = "auto"              # "auto", "cpu", "cuda:0", ...
+
+    # Output
+    save_dir: str = "saved_models"
+    model_name: str = "face_detection_yolo"
+    save: bool = True
+
+    # sample / webcam tasks
+    checkpoint: str = ""              # empty = <save_dir>/<model_name>/<model_name>
+    threshold: float = 0.95           # confidence needed to draw a box
+    camera: int = 0                   # webcam index

api/src/yolo/dataset.py

@@ -1,24 +1,14 @@
-import math
 import os
 import cv2
 import torch
-from PIL import Image
 from torch.utils.data import Dataset
 from pycocotools.coco import COCO 
 import numpy as np
 from torchvision import transforms
 import albumentations as A
 
+from src.common.boxes import image_to_cell_coords, is_center_in_grid_cell
 
-def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
-        cell_w = img_w / S
-        cell_h = img_h / S
-
-        # find which cell the center belongs to
-        gt_cell_i = min(int(x / cell_w), S - 1)
-        gt_cell_j = min(int(y / cell_h), S - 1)
-
-        return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
 
 def one_hot(index, num_classes):
     if index < 0 or index >= num_classes:
@@ -28,18 +18,6 @@ def one_hot(index, num_classes):
     encoding[index] = 1
     return encoding
 
-def turn_grid_centered(x, y, img_w, img_h, S, cell_i, cell_j):
-    cell_w = img_w / S
-    cell_h = img_h / S
-
-    cell_border_w = cell_w * cell_i
-    cell_border_h = cell_h * cell_j
-
-    return x - cell_border_w, y - cell_border_h
-
-def flip_bbox_horizontal(box, image_width):
-    x, y, w, h = box
-    return (image_width - x - w, y, w, h)
 
 class YoloDataset(Dataset):
     def __init__(self, image_dir, annotation_path, img_size=64, grid = 9, transform=False):
@@ -135,7 +113,7 @@ class YoloDataset(Dataset):
                                               img_h=self.img_size, S=self.grid, cell_i=x, cell_j=y):
                         
                         class_one_hot = one_hot(int(label), self.num_classes)
-                        x_grid_centered, y_grid_centered = turn_grid_centered(x=x_center, y=y_center, img_w=self.img_size, 
+                        x_grid_centered, y_grid_centered = image_to_cell_coords(x=x_center, y=y_center, img_w=self.img_size, 
                                               img_h=self.img_size, S=self.grid, cell_i=x, cell_j=y)
 
 

api/src/yolo/tasks.py

@@ -0,0 +1,149 @@
+"""YOLO face detector tasks. Run with: python -m src.main yolo.<task>"""
+
+import os
+import time
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision.utils import draw_bounding_boxes
+
+from src.common.checkpoints import checkpoint_path, load_checkpoint
+from src.common.device import get_device
+from src.common.training import benchmark_num_workers, fit
+from src.common.visualize import show_image
+from src.models.yolo import Yolo_model, convert_prediction
+from src.yolo.config import Config
+from src.yolo.dataset import YoloDataset
+from src.yolo.loss import YoloLoss
+
+
+def build_dataset(cfg: Config, folder: str, augment: bool):
+    return YoloDataset(
+        image_dir=folder,
+        annotation_path=os.path.join(folder, cfg.annotation_file),
+        img_size=cfg.image_size,
+        transform=augment,
+        grid=cfg.grid,
+    )
+
+
+def build_model(cfg: Config):
+    return Yolo_model(c_in=3, boxes=cfg.boxes, grid=cfg.grid, labels=cfg.labels,
+                      c_hidden=cfg.c_hidden, size_activation=cfg.size_activation)
+
+
+def load_trained(cfg: Config, device):
+    path = cfg.checkpoint or checkpoint_path(cfg.save_dir, cfg.model_name)
+    return load_checkpoint(build_model(cfg), path, device)
+
+
+def draw_prediction(label, image, threshold):
+    """Red: detected boxes, blue: cells with an object, green: empty cells."""
+    boxes, cells_obj, cells_noobj = convert_prediction(label, image, threshold)
+    if len(boxes) == 0:
+        print("no boxes")
+    image = draw_bounding_boxes(image, torch.tensor(cells_noobj), colors=(0, 255, 0))
+    image = draw_bounding_boxes(image, torch.tensor(cells_obj), colors=(0, 0, 255))
+    image = draw_bounding_boxes(image, torch.tensor(boxes), colors=(255, 0, 0))
+    show_image(image)
+
+
+def train(cfg: Config):
+    """Train the detector; saves the model with the lowest validation loss."""
+    device = get_device(cfg.device)
+    train_loader = DataLoader(build_dataset(cfg, cfg.train_dir, cfg.augment_train),
+                              batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.train_workers)
+    val_loader = DataLoader(build_dataset(cfg, cfg.val_dir, cfg.augment_val),
+                            batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.val_workers)
+
+    model = build_model(cfg).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
+
+    fit(model, YoloLoss(), optimizer, train_loader, val_loader, device,
+        epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.model_name, save=cfg.save,
+        monitor="loss")
+
+
+def view_data(cfg: Config):
+    """Show validation images with their ground-truth boxes, one at a time."""
+    dataset = build_dataset(cfg, cfg.val_dir, cfg.augment_val)
+    for images, labels in DataLoader(dataset, batch_size=1, shuffle=True):
+        draw_prediction(labels[0], images[0], cfg.threshold)
+
+
+def sample(cfg: Config):
+    """Show the saved model's predictions on validation images."""
+    device = get_device(cfg.device)
+    model = load_trained(cfg, device)
+    dataset = build_dataset(cfg, cfg.val_dir, cfg.augment_val)
+
+    for images, _ in DataLoader(dataset, batch_size=cfg.batch_size, shuffle=False):
+        with torch.inference_mode():
+            start = time.perf_counter()
+            predictions = model(images.to(device)).cpu()
+            print(f"Inference: {time.perf_counter() - start:.6f} seconds")
+        for image, prediction in zip(images, predictions):
+            draw_prediction(prediction, image, cfg.threshold)
+
+
+def benchmark_workers(cfg: Config):
+    """Time DataLoader num_workers values to pick the fastest."""
+    device = get_device(cfg.device)
+    dataset = build_dataset(cfg, cfg.train_dir, cfg.augment_train)
+    benchmark_num_workers(dataset, build_model(cfg).to(device), device, batch_size=cfg.batch_size)
+
+
+def webcam(cfg: Config):
+    """Live detection on the webcam (--camera). Press q to quit."""
+    import cv2
+
+    device = get_device(cfg.device)
+    model = load_trained(cfg, device)
+
+    cap = cv2.VideoCapture(cfg.camera)
+    if not cap.isOpened():
+        raise RuntimeError("Could not open webcam")
+
+    def draw(frame, boxes, scale_w, scale_h, color, thickness):
+        for b in boxes:
+            cv2.rectangle(frame, (int(b[0] * scale_w), int(b[1] * scale_h)),
+                          (int(b[2] * scale_w), int(b[3] * scale_h)), color, thickness)
+
+    try:
+        while True:
+            ok, frame = cap.read()
+            if not ok:
+                break
+
+            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            image = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0
+            _, H, W = image.shape
+            scale_w, scale_h = W / cfg.image_size, H / cfg.image_size
+
+            image = F.interpolate(image.unsqueeze(0), size=(cfg.image_size, cfg.image_size),
+                                  mode="bilinear", align_corners=True)
+            with torch.inference_mode():
+                prediction = model(image.to(device)).cpu()
+
+            boxes, cells_obj, cells_noobj = convert_prediction(prediction.squeeze(0), image.squeeze(0),
+                                                               threshold=cfg.threshold)
+            draw(frame, cells_noobj, scale_w, scale_h, (0, 255, 0), 1)
+            draw(frame, cells_obj, scale_w, scale_h, (255, 0, 0), 1)
+            draw(frame, boxes, scale_w, scale_h, (0, 0, 255), 4)
+
+            cv2.imshow("Webcam", frame)
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                break
+    finally:
+        cap.release()
+        cv2.destroyAllWindows()
+
+
+TASKS = {
+    "train": train,
+    "view-data": view_data,
+    "sample": sample,
+    "benchmark-workers": benchmark_workers,
+    "webcam": webcam,
+}

api/src/yolo/train_yolo_faces.py

@@ -1,160 +0,0 @@
-import os
-
-import cv2
-from tqdm import tqdm
-import torch
-import torch.nn.functional as F
-from torch.utils.data import DataLoader, random_split
-from torchvision import datasets
-from torchvision.utils import draw_bounding_boxes
-
-from src.util import TransformedSubset, test_workers_speed, visualizeImage
-from src.production.yolo_model_production import convert_prediction
-from src.yolo.yolo_dataset import YoloDataset
-from src.yolo.yolo_loss import YoloLoss
-from src.yolo.yolo_model import Yolo_model, sample, train
-
-def view_data(dataset):
-    dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=True)
-    for images, labels in iter(dataloader):
-        for batch in range(images.shape[0]):
-            image = images[batch]
-            label = labels[batch]
-            visualize_boxes(label, image)
-
-def sample_data(dataloader, model, device, SAVE_PATH):
-    for images, labels in iter(dataloader):
-        predictions = sample(model, images, device, SAVE_PATH)
-        for batch in range(predictions.shape[0]):
-            image = images[batch]
-            prediction = predictions[batch]           
-            visualize_boxes(prediction, image)
-
-def visualize_boxes(label, image, threshold=0.95):
-    boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj = convert_prediction(label, image, threshold)
-    if len(boxes_to_draw) == 0:
-        print("no labels")
-    boxes_to_draw = torch.tensor(boxes_to_draw)
-    grids_to_draw_noobj = torch.tensor(grids_to_draw_noobj)
-    grids_to_draw_obj = torch.tensor(grids_to_draw_obj)
-    image = draw_bounding_boxes(image, grids_to_draw_noobj, colors=(0, 255, 0))
-    image = draw_bounding_boxes(image, grids_to_draw_obj, colors=(0, 0, 255))
-    image = draw_bounding_boxes(image, boxes_to_draw, colors=(255, 0, 0))
-    visualizeImage(image)
-
-def use_webcam(grid, img_size):
-    # 0 = default webcam
-    cap = cv2.VideoCapture(0)
-
-    if not cap.isOpened():
-        raise Exception("Could not open webcam")
-
-    model = Yolo_model(c_in=3, boxes=1, grid=grid, labels=1)
-    state_dict = torch.load(os.path.join("saved_models", "face_detection_yolo", "face_detection_yolo"), weights_only=False)
-    model.load_state_dict(state_dict)
-    model.eval()
-
-    while True:
-        ret, frame = cap.read()
-
-        if not ret:
-            break
-
-        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-
-        image = torch.from_numpy(rgb).permute(2, 0, 1).float()
-        C, H, W = image.shape
-        scale_w = W / img_size
-        scale_h = H / img_size
-
-        image = image / 255.0
-        
-        image = image.unsqueeze(0)
-        image = F.interpolate(image, size=(img_size, img_size), mode="bilinear", align_corners=True) #different modes?
-
-        prediction = model(image)
-
-        bboxes, grid_ob, grid_noob = convert_prediction(prediction.squeeze(0), image.squeeze(0), threshold=0.9)
-
-        for bbox in grid_noob:
-            xmin = int(bbox[0] * scale_w)
-            ymin = int(bbox[1] * scale_h)
-            xmax = int(bbox[2] * scale_w)
-            ymax = int(bbox[3] * scale_h)
-            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 0), 1)
-
-        for bbox in grid_ob:
-            xmin = int(bbox[0] * scale_w)
-            ymin = int(bbox[1] * scale_h)
-            xmax = int(bbox[2] * scale_w)
-            ymax = int(bbox[3] * scale_h)
-            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255, 0, 0), 1)
-
-        for bbox in bboxes:
-            xmin = int(bbox[0] * scale_w)
-            ymin = int(bbox[1] * scale_h)
-            xmax = int(bbox[2] * scale_w)
-            ymax = int(bbox[3] * scale_h)
-            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 0, 255), 4)
-
-        cv2.imshow("Webcam", frame)
-        if cv2.waitKey(1) & 0xFF == ord("q"):
-            break
-
-    cap.release()
-    cv2.destroyAllWindows()
-
-
-def train_yolo():
-    SAVE_PATH = "./saved_models"
-    IMAGE_SIZE = 64
-    GRID = 6
-    BATCH_SIZE = 64
-
-    dataset = YoloDataset(
-        image_dir="data/faces_scenery/train", 
-        annotation_path="data/faces_scenery/train/_annotations.coco.json",
-        img_size=IMAGE_SIZE,
-        transform=True,
-        grid=GRID
-    )
-
-    dataset_valid = YoloDataset(
-        image_dir="data/faces_scenery/test", 
-        annotation_path="data/faces_scenery/test/_annotations.coco.json",
-        img_size=IMAGE_SIZE,
-        transform=True,
-        grid=GRID
-    )
-
-
-    train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
-    val_loader = DataLoader(dataset_valid, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
-
-    device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
-    print("Using device", device)
-
-    model = Yolo_model(c_in=3, boxes=1, grid=GRID, labels=1)
-    model.to(device)
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
-    loss_module = YoloLoss()
-
-
-    view_data(dataset_valid)
-    exit()
-
-
-    #sample_data(dataloader=val_loader, model=model, device=device, SAVE_PATH=SAVE_PATH)
-    #exit()
-
-
-    #test_workers_speed(dataset, model)
-    #exit()
-
-
-    #use_webcam(GRID, IMAGE_SIZE)
-    #exit()
-
-
-    train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, 
-                optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, model_name="face_detection_yolo")

api/src/yolo/yolo_model.py

@@ -1,162 +0,0 @@
-import os
-import time
-import torch
-from tqdm import tqdm
-import torch.nn as nn
-import torch.nn.functional as F
-
-class Yolo_Conv_Block(nn.Module):
-    def __init__(self, c_in, c_hidden, c_out, kernel_size):
-        super().__init__()
-
-        self.model = nn.Sequential(
-            nn.Conv2d(in_channels=c_in, out_channels=c_hidden, kernel_size=kernel_size, padding=kernel_size//2),
-            nn.BatchNorm2d(c_hidden),
-            nn.LeakyReLU(),
-            nn.Conv2d(in_channels=c_hidden, out_channels=c_out, kernel_size=1)
-        )
-
-    def forward(self, x):
-        return self.model(x)
-    
-class SkipBlock(nn.Module):
-    def __init__(self, c_in, c_out, kernel_size=3):
-        super().__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size//2),
-            nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
-            nn.LeakyReLU(inplace=True),
-
-            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
-            nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
-            nn.LeakyReLU(inplace=True),
-
-            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
-            nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
-            nn.LeakyReLU(inplace=True)
-        )
-        self.conv_skip = nn.Sequential(
-            nn.Conv2d(c_in, c_out, 1),
-            nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
-            nn.LeakyReLU(inplace=True)
-        )
-
-    def forward(self, x):
-        return(F.dropout(F.leaky_relu(self.conv_skip(x) + self.conv(x), inplace=True), p=0.3))
-
-
-class Yolo_model(nn.Module):
-    def __init__(self, c_in, boxes, grid, labels, c_hidden=16):
-        super().__init__()
-
-        self.model = nn.Sequential(
-            nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups=c_hidden//8, num_channels=c_hidden),
-            nn.LeakyReLU(inplace=True),
-
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            SkipBlock(c_in=c_hidden, c_out=c_hidden),
-            
-            SkipBlock(c_in=c_hidden, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
-
-            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*4),
-            SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
-            SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
-            SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
-
-            nn.Conv2d(c_hidden*4, c_hidden*8, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups=c_hidden//2, num_channels=c_hidden*8),
-            nn.LeakyReLU(inplace=True),
-            nn.Dropout(0.3),
-
-            nn.AdaptiveAvgPool2d((grid, grid)),
-            nn.Conv2d(c_hidden*8, boxes*5 + labels, kernel_size=1)
-        )
-
-    def forward(self, x):
-        x = self.model(x).permute(0, 2, 3, 1)
-        return F.sigmoid(x)
-    
-
-def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True):
-    best_val = torch.finfo(torch.float32).max
-    device = next(model.parameters()).device
-
-    for epoch in range(500):
-        ############
-        # Training #
-        ############
-        model.train()
-
-        count, lossCount = 0, 0.
-        for images, labels in tqdm(train_loader, desc=f"Train", leave=False):
-            images = images.to(device)
-            labels = labels.to(device)
-
-            prediction = model(images)
-
-            loss = loss_module(prediction, labels)
-            lossCount += loss.item()
-
-            optimizer.zero_grad()
-
-            loss.backward()
-
-            optimizer.step()
-
-            count += images.size(0)
-
-        train_loss = lossCount / count
-
-        torch.cuda.empty_cache()
-
-        ##############
-        # Validation #
-        ##############
-        model.eval()
-
-        count, lossCount = 0, 0.
-        for images, labels in tqdm(val_loader, desc=f"Test", leave=False):
-            with torch.no_grad():
-                images = images.to(device)
-                labels = labels.to(device)
-
-                prediction = model(images)
-
-                loss = loss_module(prediction, labels)
-                lossCount += loss.item()
-
-                count += images.size(0)
-
-        val_loss = lossCount / count
-
-        if(saving and best_val > val_loss):
-                best_val = val_loss
-                save_dir = os.path.join(SAVE_PATH, model_name)
-                os.makedirs(save_dir, exist_ok=True)
-
-                save_path = os.path.join(save_dir, model_name)
-                torch.save(model.state_dict(), save_path)
-
-        print(f"epoch: {epoch+1} | train loss: {int(train_loss * 100000) / 100} | val loss: {int(val_loss * 100000) / 100}")
-        torch.cuda.empty_cache()
-    return best_val
-
-
-def sample(model, img, device, SAVE_PATH, model_name="face_detection_yolo", folder="face_detection_yolo"):
-    with torch.no_grad():
-        full_path = os.path.join(SAVE_PATH, folder, model_name)
-        state_dict = torch.load(full_path, weights_only=False)
-        model.load_state_dict(state_dict)
-        model.eval()
-        img = img.to(device)
-        start = time.perf_counter()
-        pred = model(img)
-        end = time.perf_counter()
-        print(f"Elapsed: {end - start:.6f} seconds") 
-        return pred

api/start_server.bat

@@ -0,0 +1 @@
+python -m src.main server.run

api/start_server.sh

@@ -0,0 +1 @@
+uvicorn src.production.app:app --host 0.0.0.0 --port 8000 --ws-max-size 1048576