diff --git a/api/src/main.py b/api/src/main.py new file mode 100644 index 0000000..f3787f0 --- /dev/null +++ b/api/src/main.py @@ -0,0 +1,4 @@ +from yolo.train_yolo_faces import train_yolo + + +train_yolo() \ No newline at end of file diff --git a/api/src/r_cnn/cocoDetectionDataset.py b/api/src/r_cnn/cocoDetectionDataset.py index eb4b3ab..27b76bf 100644 --- a/api/src/r_cnn/cocoDetectionDataset.py +++ b/api/src/r_cnn/cocoDetectionDataset.py @@ -57,69 +57,4 @@ class CocoDetectionDataset(Dataset): if self.transforms: image = self.transforms(image) - return image, target - - -class CocoDetectionDatasetResized(Dataset): - def __init__(self, image_dir, annotation_path, img_size=64, transforms=None): - self.image_dir = image_dir - self.coco = COCO(annotation_path) - self.image_ids = list(self.coco.imgs.keys()) - self.transforms = transforms - self.img_size = img_size - - def __len__(self): - return len(self.image_ids) - - def __getitem__(self, idx): - image_id = self.image_ids[idx] - image_info = self.coco.loadImgs(image_id)[0] - image_path = os.path.join(self.image_dir, image_info['file_name']) - - image = Image.open(image_path).convert("RGB") - - orig_w, orig_h = image.size - - scale_w = self.img_size / orig_w - scale_h = self.img_size / orig_h - - # Load annotations - annotation_ids = self.coco.getAnnIds(imgIds=image_id) - annotations = self.coco.loadAnns(annotation_ids) - - boxes = [] - labels = [] - - for obj in annotations: - xmin, ymin, width, height = obj['bbox'] - xmin, ymin, width, height = float(xmin), float(ymin), float(width), float(height) - - xmin = xmin * scale_w - ymin = ymin * scale_h - xmax = (xmin + width * scale_w) - ymax = (ymin + height * scale_h) - - boxes.append([xmin, ymin, xmax, ymax]) - labels.append(obj['category_id']) - - boxes = torch.tensor(boxes, dtype=torch.float32) - labels = torch.tensor(labels, dtype=torch.long) - - area = torch.tensor([obj['area'] * scale_w * scale_h for obj in annotations], dtype=torch.float32) - iscrowd = torch.tensor([obj.get('iscrowd', 0) for obj in annotations], dtype=torch.long) - - target = { - "boxes": boxes, - "labels": labels, - "image_id": torch.tensor([image_id]), - "area": area, - "iscrowd": iscrowd - } - - # resize image - image = image.resize((self.img_size, self.img_size)) - - if self.transforms: - image = self.transforms(image) - return image, target \ No newline at end of file diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py index 5f49384..bf74110 100644 --- a/api/src/yolo/train_yolo_faces.py +++ b/api/src/yolo/train_yolo_faces.py @@ -2,9 +2,9 @@ import os from torchvision import transforms from tqdm import tqdm -from yolo_model import train, Yolo_model -from yolo_loss import YoloLoss -from cocoDetectionDataset import CocoDetectionDatasetResized +from yolo.yolo_dataset import YoloDataset +from yolo.yolo_model import train, Yolo_model +from yolo.yolo_loss import YoloLoss from util import TransformedSubset, visualizeImage from torch.utils.data import DataLoader, random_split from torchvision import datasets @@ -14,42 +14,41 @@ from torchvision.io import read_image from torch.utils.data import WeightedRandomSampler +def train_yolo(): + SAVE_PATH = "./saved_models" + IMAGE_SIZE = 64 + GRID = 9 + BATCH_SIZE = 256 -SAVE_PATH = "./saved_models" -IMAGE_SIZE = 128 -GRID = 9 -BATCH_SIZE = 1 + transform = transforms.Compose([ + transforms.ToTensor() + ]) -transform = transforms.Compose([ - transforms.ToTensor() -]) + dataset = YoloDataset( + image_dir="data/faces/train", + annotation_path="data/faces/train/_annotations.coco.json", + img_size=IMAGE_SIZE, + transforms=transform + ) -dataset = CocoDetectionDatasetResized( - image_dir="data/faces/train", - annotation_path="data/faces/train/_annotations.coco.json", - img_size=IMAGE_SIZE, - transforms=transform -) + train_size = int(0.8 * len(dataset)) + val_size = len(dataset) - train_size -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_subset, val_subset = random_split(dataset, [train_size, val_size]) -train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=lambda x: tuple(zip(*x))) -val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=lambda x: tuple(zip(*x))) + train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=False) + val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False) -device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0") -print("Using device", device) + 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, c_hidden=32, boxes=2, img_size=IMAGE_SIZE, grid=GRID, labels=1) -model.to(device) -optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4) -loss_module = YoloLoss() + model = Yolo_model(c_in=3, c_hidden=32, boxes=1, img_size=IMAGE_SIZE, grid=GRID, labels=1) + model.to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4) + loss_module = YoloLoss() -train(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_yolo", img_size=IMAGE_SIZE, num_classes=1, grid=GRID) -exit() \ No newline at end of file + train(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_yolo") \ No newline at end of file diff --git a/api/src/yolo/yolo_dataset.py b/api/src/yolo/yolo_dataset.py new file mode 100644 index 0000000..bbf4b53 --- /dev/null +++ b/api/src/yolo/yolo_dataset.py @@ -0,0 +1,105 @@ +import os +import torch +from PIL import Image +from torch.utils.data import Dataset +from pycocotools.coco import COCO +import numpy as np + + +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 = int(y / cell_h) + gt_cell_j = int(x / cell_w) + + 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: + raise ValueError(f"Index out of range, index: {index}, num_classes: {num_classes}") + + encoding = [0] * num_classes + encoding[index] = 1 + return encoding + +class YoloDataset(Dataset): + def __init__(self, image_dir, annotation_path, img_size=64, grid = 9, transforms=None): + self.image_dir = image_dir + self.coco = COCO(annotation_path) + self.image_ids = list(self.coco.imgs.keys()) + self.transforms = transforms + self.img_size = img_size + self.grid = grid + + def __len__(self): + return len(self.image_ids) + + def __getitem__(self, idx): + image_id = self.image_ids[idx] + image_info = self.coco.loadImgs(image_id)[0] + image_path = os.path.join(self.image_dir, image_info['file_name']) + + image = Image.open(image_path).convert("RGBA").convert("RGB") + + orig_w, orig_h = image.size + + scale_w = self.img_size / orig_w + scale_h = self.img_size / orig_h + + # Load annotations + annotation_ids = self.coco.getAnnIds(imgIds=image_id) + annotations = self.coco.loadAnns(annotation_ids) + + boxes = [] + labels = [] + + for obj in annotations: + xmin, ymin, width, height = obj['bbox'] + xmin, ymin, width, height = float(xmin), float(ymin), float(width), float(height) + + xmin = xmin * scale_w + ymin = ymin * scale_h + xmax = (xmin + width * scale_w) + ymax = (ymin + height * scale_h) + + boxes.append([xmin, ymin, xmax, ymax]) + labels.append(obj['category_id'] - 1) + + ground_truth = list(zip(boxes, labels)) + num_labels = max(labels) + 1 if len(labels) > 0 else 1 + #[S, S, (x+y+w+h+c+C)] + targets = np.zeros((self.grid, self.grid, 5 + num_labels)) + for x in range(self.grid): + for y in range(self.grid): + for i in range(len(ground_truth)): + box = ground_truth[i][0] + label = ground_truth[i][1] + + if is_center_in_grid_cell(x=box[0]+box[2]/2, y=box[1]+box[3]/2, img_w=self.img_size, + img_h=self.img_size, S=self.grid, cell_i=x, cell_j=y): + + class_one_hot = one_hot(label, num_labels) + targets[x, y, 0] = box[0] + targets[x, y, 1] = box[1] + targets[x, y, 2] = box[2] + targets[x, y, 3] = box[3] + targets[x, y, 4] = 1 + for i in range(len(class_one_hot)): + targets[x, y, i + 5] = class_one_hot[i] + + del ground_truth[i] + + break + + + targets = torch.tensor(targets, dtype=torch.float32) + + # resize image + image = image.resize((self.img_size, self.img_size)) + + if self.transforms: + image = self.transforms(image) + + return image, targets \ No newline at end of file diff --git a/api/src/yolo/yolo_loss.py b/api/src/yolo/yolo_loss.py index 0d3bd18..a53f108 100644 --- a/api/src/yolo/yolo_loss.py +++ b/api/src/yolo/yolo_loss.py @@ -5,103 +5,20 @@ class YoloLoss(nn.Module): def __init__(self): super(YoloLoss, self).__init__() - @staticmethod - 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 = int(y / cell_h) - gt_cell_j = int(x / cell_w) - - return (gt_cell_i == cell_i) and (gt_cell_j == cell_j) - - @staticmethod - def save_sqrt(i): - return torch.sqrt(torch.clamp(i, min=1e-6)) - - @staticmethod - def iou(boxA, boxB): - xA = torch.max(boxA[0], boxB[0]) - yA = torch.max(boxA[1], boxB[1]) - xB = torch.min(boxA[2], boxB[2]) - yB = torch.min(boxA[3], boxB[3]) - - inter_area = torch.clamp(xB - xA, min=0) * torch.clamp(yB - yA, min=0) - - 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 + 1e-6) - - @staticmethod - def xywh_to_xyxy(box): - x, y, w, h = box - - return torch.stack([ - x - w / 2, - y - h / 2, - x + w / 2, - y + h / 2 - ]) - - def forward(self, pred, target, labels, cell_size, img_w, img_h, factor_no_object = 0.5, factor_object = 5): - #pred: [SxSx(B(x,y,w,h,con)C)] - #target: [B(x, y, w, h, C)] + def forward(self, predictions, targets, lambda_coord=5, lambda_noobj=0.5): + pred_boxes = predictions[..., :4] + pred_conf = predictions[..., 4] + pred_classes = predictions[..., 5:] + target_boxes = targets[..., :4] + target_conf = targets[..., 4] + target_classes = targets[..., 5:] - for s0 in range(pred.shape[0]): - for s1 in range(pred.shape[1]): - boxes = pred[s0, s1, :(-1*labels)] - pred_label = pred[s0, s1, -1:] + box_loss = lambda_coord * torch.mean((pred_boxes - target_boxes) ** 2) - ground_truth_box = None - for b in range(target.shape[0] // (4 + labels)): - box = target[(b)*(4 + labels):(b+1)*(4 + labels)] - if self.is_center_in_grid_cell(box[0]+box[2]/2, box[1]+box[3]/2, img_w, img_h, cell_size, s0, s1): - ground_truth_box = box[:4] - ground_truth_label = box[4:] - break + obj_loss = torch.mean((pred_conf[target_conf == 1] - target_conf[target_conf == 1]) ** 2) + noobj_loss = lambda_noobj * torch.mean((pred_conf[target_conf == 0]) ** 2) + class_loss = torch.mean((pred_classes[target_conf == 1] - target_classes[target_conf == 1]) ** 2) - coord_loss = torch.tensor(0., device=pred.device) - size_loss = torch.tensor(0., device=pred.device) - prob_loss = torch.tensor(0., device=pred.device) - label_loss = torch.tensor(0., device=pred.device) - - if ground_truth_box is not None: - - iou_stats = [] - - for b in range(boxes.shape[0] // 5): - box = boxes[(b)*5:(b+1)*5] - - iou_stats.append(self.iou(self.xywh_to_xyxy(box[0:4]), self.xywh_to_xyxy(ground_truth_box[0:4]))) - iou_tensor = torch.stack(iou_stats) - responsible = torch.argmax(iou_tensor) - - for b in range(boxes.shape[0] // 5): - box = boxes[(b)*5:(b+1)*5] - - is_responsible = (b == responsible).float() - - coord_loss += factor_object * is_responsible * ((box[0] - ground_truth_box[0])**2 + (box[1] - ground_truth_box[1])**2) - size_loss += factor_object * is_responsible * ((self.save_sqrt(box[2]) - self.save_sqrt(ground_truth_box[2]))**2 + - (self.save_sqrt(box[3]) - self.save_sqrt(ground_truth_box[3]))**2) - - prob_loss += is_responsible * (box[4] - 1)**2 - prob_loss += factor_no_object * (1 - is_responsible) * box[4]**2 - - - for i in range(pred_label.shape[0]): - label_loss += ((pred_label[i] - ground_truth_label[i])**2) / pred_label.shape[0] - - - else: - for b in range(boxes.shape[0] // 5): - box = boxes[(b)*5:(b+1)*5] - - prob_loss += factor_no_object * box[4] ** 2 - - return (label_loss + prob_loss + coord_loss + size_loss) / pred.shape[0]**2 + total_loss = box_loss + obj_loss + noobj_loss + class_loss + return total_loss diff --git a/api/src/yolo/yolo_model.py b/api/src/yolo/yolo_model.py index d95a8e7..cdbbf0c 100644 --- a/api/src/yolo/yolo_model.py +++ b/api/src/yolo/yolo_model.py @@ -1,11 +1,7 @@ import os -import cv2 import torch from tqdm import tqdm import torch.nn as nn -import torch.nn.functional as F -import torchvision.transforms.functional as TF -import time class Yolo_model(nn.Module): @@ -16,16 +12,22 @@ class Yolo_model(nn.Module): self.model = nn.Sequential( nn.Conv2d(in_channels=c_in, out_channels=c_hidden, kernel_size=7, padding=3), + nn.BatchNorm2d(c_hidden), nn.LeakyReLU(), nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(c_hidden), nn.LeakyReLU(), nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(c_hidden), nn.LeakyReLU(), nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(c_hidden), nn.LeakyReLU(), nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(c_hidden), nn.LeakyReLU(), nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(c_hidden), nn.LeakyReLU(), nn.Flatten(), nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels)) @@ -37,29 +39,8 @@ class Yolo_model(nn.Module): x = x.reshape(batch, self.grid, self.grid, -1) return x -def create_stack(images, annotations, num_classes, device): - data = [] - labels = [] - for i in range(len(images)): - image = images[i] - annotation = annotations[i] - - label_train = [] - for box, label in tuple(zip(annotation["boxes"], annotation["labels"])): - box = box - label = 0 - one_hot = F.one_hot(torch.tensor(label, dtype=torch.long), num_classes=num_classes) - label_train.append(torch.cat((box, one_hot), dim=0)) - - label_torch = torch.stack(label_train).reshape(-1).to(device) - labels.append(label_torch) - data.append(image) - - return data, labels - - -def train(model, num_classes, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, grid, saving=True, device="cpu", img_size=64): +def train(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): @@ -69,18 +50,13 @@ def train(model, num_classes, loss_module, train_loader, val_loader, optimizer, model.train() count, lossCount = 0, 0. - for images, annotations in tqdm(train_loader, desc=f"Train", leave=False): - data, labels = create_stack(images, annotations, num_classes, device) + for images, labels in tqdm(train_loader, desc=f"Train", leave=False): + images = images.to(device) + labels = labels.to(device) - data = torch.stack(data).to(device) + prediction = model(images) - prediction = model(data) - - loss = [] - for pred, label in tuple(zip(prediction, labels)): - loss.append(loss_module(pred = pred, target= label, labels=num_classes, cell_size=grid, img_w = img_size, img_h = img_size)) - loss = torch.stack(loss).to(device) - loss = loss.sum() + loss = loss_module(prediction, labels) lossCount += loss.item() optimizer.zero_grad() @@ -89,7 +65,7 @@ def train(model, num_classes, loss_module, train_loader, val_loader, optimizer, optimizer.step() - count += data.size(0) + count += images.size(0) train_loss = lossCount / count @@ -101,23 +77,20 @@ def train(model, num_classes, loss_module, train_loader, val_loader, optimizer, model.eval() count, lossCount = 0, 0. - for images, annotations in tqdm(val_loader, desc=f"Test", leave=False): + for images, labels in tqdm(val_loader, desc=f"Test", leave=False): with torch.no_grad(): - data, labels = create_stack(images, annotations, num_classes, device) + images = images.to(device) + labels = labels.to(device) - data = torch.stack(data).to(device) + prediction = model(images) - prediction = model(data) - - loss = [] - for pred, label in tuple(zip(prediction, labels)): - loss.append(loss_module(pred = pred, target= label, labels=num_classes, cell_size=grid, img_w = img_size, img_h = img_size)) - loss = torch.stack(loss).to(device) + loss = loss_module(prediction, labels) + lossCount += loss.item() - lossCount += loss.sum().item() + lossCount += loss.sum().item() - count += data.size(0) + count += images.size(0) val_loss = lossCount / count