diff --git a/api/.dockerignore b/api/.dockerignore index c0d77c9..eda9ba6 100644 --- a/api/.dockerignore +++ b/api/.dockerignore @@ -1,9 +1,13 @@ -# Allowlist: the Dockerfile only copies these three things, so nothing -# else (training code, datasets, saved models, caches) is sent to the daemon. +# Allowlist: only what the server needs is sent to the Docker daemon, so +# 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__ diff --git a/api/Dockerfile b/api/Dockerfile index 051f95a..1567f58 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -12,9 +12,12 @@ 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 -ENTRYPOINT ["bash", "./start_server.sh"] \ No newline at end of file +ENTRYPOINT ["bash", "./start_server.sh"] diff --git a/api/src/bird_cnn/bird_cnn.py b/api/src/bird_cnn/bird_cnn.py deleted file mode 100644 index 1873c9f..0000000 --- a/api/src/bird_cnn/bird_cnn.py +++ /dev/null @@ -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)) \ No newline at end of file diff --git a/api/src/bird_cnn/config.py b/api/src/bird_cnn/config.py new file mode 100644 index 0000000..17b439b --- /dev/null +++ b/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 = // diff --git a/api/src/bird_cnn/tasks.py b/api/src/bird_cnn/tasks.py new file mode 100644 index 0000000..48dccfe --- /dev/null +++ b/api/src/bird_cnn/tasks.py @@ -0,0 +1,87 @@ +"""Bird classifier tasks. Run with: python -m src.main bird.""" + +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, +} diff --git a/api/src/bird_cnn/train_bird.py b/api/src/bird_cnn/train_bird.py deleted file mode 100644 index d0bc2a9..0000000 --- a/api/src/bird_cnn/train_bird.py +++ /dev/null @@ -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}") \ No newline at end of file diff --git a/api/src/common/__init__.py b/api/src/common/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/src/common/boxes.py b/api/src/common/boxes.py new file mode 100644 index 0000000..4d4e952 --- /dev/null +++ b/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)] diff --git a/api/src/common/checkpoints.py b/api/src/common/checkpoints.py new file mode 100644 index 0000000..de77aea --- /dev/null +++ b/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: // + 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() diff --git a/api/src/common/data.py b/api/src/common/data.py new file mode 100644 index 0000000..1825a82 --- /dev/null +++ b/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)) diff --git a/api/src/common/device.py b/api/src/common/device.py new file mode 100644 index 0000000..f003057 --- /dev/null +++ b/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 diff --git a/api/src/common/training.py b/api/src/common/training.py new file mode 100644 index 0000000..df24b56 --- /dev/null +++ b/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") diff --git a/api/src/common/visualize.py b/api/src/common/visualize.py new file mode 100644 index 0000000..a7358d7 --- /dev/null +++ b/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() diff --git a/api/src/main.py b/api/src/main.py index fd1a190..700b09f 100644 --- a/api/src/main.py +++ b/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//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(): - train_yolo() +def load_group(group: str): + 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() \ No newline at end of file + main() diff --git a/api/src/models/__init__.py b/api/src/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/src/production/bird_cnn_production.py b/api/src/models/bird_cnn.py similarity index 58% rename from api/src/production/bird_cnn_production.py rename to api/src/models/bird_cnn.py index 27d240c..3d3a9c4 100644 --- a/api/src/production/bird_cnn_production.py +++ b/api/src/models/bird_cnn.py @@ -1,36 +1,34 @@ +"""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): - x = self.depthwise(x) - x = self.bn1(x) - x = F.relu(x) +class BirdSpecies(Enum): + """Classes of the deployed 7-species model (build_models/bird_cnn).""" + Common_Kingfisher = 0 + 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): super().__init__() self.conv = nn.Sequential( - nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size//2), + 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.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.Conv2d(c_out, c_out, kernel_size, padding=kernel_size // 2), nn.BatchNorm2d(c_out), nn.ReLU(inplace=True) ) @@ -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): @@ -57,20 +55,20 @@ class Bird_CNN(nn.Module): 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), + 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.Linear(c_hidden * 4, c_out), nn.Dropout(0.3) ) def forward(self, x): - return self.model(x) \ No newline at end of file + return self.model(x) diff --git a/api/src/models/r_cnn.py b/api/src/models/r_cnn.py new file mode 100644 index 0000000..7fde45f --- /dev/null +++ b/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 diff --git a/api/src/models/yolo.py b/api/src/models/yolo.py new file mode 100644 index 0000000..0ddfcbd --- /dev/null +++ b/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 diff --git a/api/src/production/__init__.py b/api/src/production/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/src/production/app.py b/api/src/production/app.py new file mode 100644 index 0000000..f691e2f --- /dev/null +++ b/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"} diff --git a/api/src/production/config.py b/api/src/production/config.py new file mode 100644 index 0000000..ceefebf --- /dev/null +++ b/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 diff --git a/api/src/production/inference.py b/api/src/production/inference.py new file mode 100644 index 0000000..8e8283e --- /dev/null +++ b/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] diff --git a/api/src/production/limits.py b/api/src/production/limits.py new file mode 100644 index 0000000..3d870b9 --- /dev/null +++ b/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) diff --git a/api/src/production/predict.py b/api/src/production/predict.py new file mode 100644 index 0000000..05dc7ea --- /dev/null +++ b/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] diff --git a/api/src/production/reviews.py b/api/src/production/reviews.py new file mode 100644 index 0000000..156c009 --- /dev/null +++ b/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"} diff --git a/api/src/production/server.py b/api/src/production/server.py deleted file mode 100644 index a491a0e..0000000 --- a/api/src/production/server.py +++ /dev/null @@ -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"} diff --git a/api/src/production/settings.py b/api/src/production/settings.py new file mode 100644 index 0000000..d5f371a --- /dev/null +++ b/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) diff --git a/api/src/production/start_server.bat b/api/src/production/start_server.bat deleted file mode 100644 index 6f52117..0000000 --- a/api/src/production/start_server.bat +++ /dev/null @@ -1 +0,0 @@ -python -m uvicorn server:app --ws-max-size 1048576 --reload \ No newline at end of file diff --git a/api/src/production/start_server.sh b/api/src/production/start_server.sh deleted file mode 100644 index 484e4d1..0000000 --- a/api/src/production/start_server.sh +++ /dev/null @@ -1 +0,0 @@ -uvicorn server:app --host 0.0.0.0 --port 8000 --ws-max-size 1048576 diff --git a/api/src/production/tasks.py b/api/src/production/tasks.py new file mode 100644 index 0000000..c20f0d0 --- /dev/null +++ b/api/src/production/tasks.py @@ -0,0 +1,16 @@ +"""API server tasks. Run with: python -m src.main server.""" + +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, +} diff --git a/api/src/production/yolo_model_production.py b/api/src/production/yolo_model_production.py deleted file mode 100644 index 82d6531..0000000 --- a/api/src/production/yolo_model_production.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/api/src/r_cnn/config.py b/api/src/r_cnn/config.py new file mode 100644 index 0000000..e50a65c --- /dev/null +++ b/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 diff --git a/api/src/r_cnn/createObjectDetectionDataset.py b/api/src/r_cnn/createObjectDetectionDataset.py deleted file mode 100644 index ee996bd..0000000 --- a/api/src/r_cnn/createObjectDetectionDataset.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/api/src/r_cnn/crops.py b/api/src/r_cnn/crops.py new file mode 100644 index 0000000..451deff --- /dev/null +++ b/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 diff --git a/api/src/r_cnn/cocoDetectionDataset.py b/api/src/r_cnn/dataset.py similarity index 100% rename from api/src/r_cnn/cocoDetectionDataset.py rename to api/src/r_cnn/dataset.py diff --git a/api/src/r_cnn/detect.py b/api/src/r_cnn/detect.py new file mode 100644 index 0000000..5e1b72b --- /dev/null +++ b/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) diff --git a/api/src/r_cnn/r_cnn.py b/api/src/r_cnn/r_cnn.py deleted file mode 100644 index f630786..0000000 --- a/api/src/r_cnn/r_cnn.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/api/src/r_cnn/r_cnn_test.py b/api/src/r_cnn/r_cnn_test.py deleted file mode 100644 index 5e81dec..0000000 --- a/api/src/r_cnn/r_cnn_test.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/api/src/r_cnn/tasks.py b/api/src/r_cnn/tasks.py new file mode 100644 index 0000000..0a6d039 --- /dev/null +++ b/api/src/r_cnn/tasks.py @@ -0,0 +1,126 @@ +"""R-CNN style experiments. Run with: python -m src.main rcnn.""" + +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, +} diff --git a/api/src/r_cnn/train_face_detection_r_cnn.py b/api/src/r_cnn/train_face_detection_r_cnn.py deleted file mode 100644 index 383e7a4..0000000 --- a/api/src/r_cnn/train_face_detection_r_cnn.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/api/src/util.py b/api/src/util.py deleted file mode 100644 index e3ab3f0..0000000 --- a/api/src/util.py +++ /dev/null @@ -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)) \ No newline at end of file diff --git a/api/src/yolo/config.py b/api/src/yolo/config.py new file mode 100644 index 0000000..a701b40 --- /dev/null +++ b/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 = // + threshold: float = 0.95 # confidence needed to draw a box + camera: int = 0 # webcam index diff --git a/api/src/yolo/yolo_dataset.py b/api/src/yolo/dataset.py similarity index 85% rename from api/src/yolo/yolo_dataset.py rename to api/src/yolo/dataset.py index 732ac3f..ca69f94 100644 --- a/api/src/yolo/yolo_dataset.py +++ b/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) diff --git a/api/src/yolo/yolo_loss.py b/api/src/yolo/loss.py similarity index 100% rename from api/src/yolo/yolo_loss.py rename to api/src/yolo/loss.py diff --git a/api/src/yolo/tasks.py b/api/src/yolo/tasks.py new file mode 100644 index 0000000..6e80fbc --- /dev/null +++ b/api/src/yolo/tasks.py @@ -0,0 +1,149 @@ +"""YOLO face detector tasks. Run with: python -m src.main yolo.""" + +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, +} diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py deleted file mode 100644 index 224577b..0000000 --- a/api/src/yolo/train_yolo_faces.py +++ /dev/null @@ -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") \ No newline at end of file diff --git a/api/src/yolo/yolo_model.py b/api/src/yolo/yolo_model.py deleted file mode 100644 index 7ce4329..0000000 --- a/api/src/yolo/yolo_model.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/api/start_server.bat b/api/start_server.bat new file mode 100644 index 0000000..eaf5421 --- /dev/null +++ b/api/start_server.bat @@ -0,0 +1 @@ +python -m src.main server.run diff --git a/api/start_server.sh b/api/start_server.sh new file mode 100644 index 0000000..70d0ceb --- /dev/null +++ b/api/start_server.sh @@ -0,0 +1 @@ +uvicorn src.production.app:app --host 0.0.0.0 --port 8000 --ws-max-size 1048576