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