diff --git a/api/src/util.py b/api/src/util.py index 127ed1b..e3ab3f0 100644 --- a/api/src/util.py +++ b/api/src/util.py @@ -64,7 +64,7 @@ def iou(boxA, boxB): 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=16,pin_memory=True) + 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): diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py index 1528360..63d9a5b 100644 --- a/api/src/yolo/train_yolo_faces.py +++ b/api/src/yolo/train_yolo_faces.py @@ -9,11 +9,7 @@ from util import TransformedSubset, test_workers_speed, visualizeImage from torch.utils.data import DataLoader, random_split from torchvision import datasets import torch -import torch.nn as nn -from torchvision.io import read_image -from torch.utils.data import WeightedRandomSampler from torchvision.utils import draw_bounding_boxes -import time import torch.nn.functional as F def xy_center_to_edges(xcenter, ycenter, width, height): @@ -36,17 +32,14 @@ def view_data(dataset): def sample_data(dataloader, model, device, SAVE_PATH): for images, labels in iter(dataloader): - start = time.perf_counter() predictions = sample(model, images, device, SAVE_PATH) - end = time.perf_counter() - print(f"Elapsed: {end - start:.6f} seconds") for batch in range(predictions.shape[0]): image = images[batch] prediction = predictions[batch] visualize_boxes(prediction, image) -def visualize_boxes(label, image): - boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj = convert_prediction(label, 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") return @@ -87,14 +80,14 @@ def convert_prediction(label, image, threshold=0.9): boxes_to_draw.append(xy_center_to_edges(boxx, boxy, boxw, boxh)) #xmin, ymin, xmax, ymax return boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj -def use_webcam(): +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=5, labels=1) + 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() @@ -109,17 +102,17 @@ def use_webcam(): image = torch.from_numpy(rgb).permute(2, 0, 1).float() C, H, W = image.shape - scale_w = W / 64 - scale_h = H / 64 + scale_w = W / img_size + scale_h = H / img_size image = image / 255.0 image = image.unsqueeze(0) - image = F.interpolate(image, size=(64, 64), mode="bilinear", align_corners=False) #different modes? + 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.995) + 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) @@ -153,8 +146,8 @@ def use_webcam(): def train_yolo(): SAVE_PATH = "./saved_models" IMAGE_SIZE = 64 - GRID = 3 - BATCH_SIZE = 32 + GRID = 6 + BATCH_SIZE = 64 dataset = YoloDataset( image_dir="data/faces_2/train", @@ -168,12 +161,12 @@ def train_yolo(): image_dir="data/faces_2/test", annotation_path="data/faces_2/test/_annotations.coco.json", img_size=IMAGE_SIZE, - transform=False, + transform=True, grid=GRID ) - train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0) + 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") @@ -189,15 +182,16 @@ def train_yolo(): #exit() - sample_data(dataloader=val_loader, model=model, device=device, SAVE_PATH=SAVE_PATH) - exit() + #sample_data(dataloader=val_loader, model=model, device=device, SAVE_PATH=SAVE_PATH) + #exit() #test_workers_speed(dataset, model) #exit() - #use_webcam() - #exit() + + use_webcam(GRID, IMAGE_SIZE) + exit() train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, diff --git a/api/src/yolo/yolo_dataset.py b/api/src/yolo/yolo_dataset.py index d824def..d368b88 100644 --- a/api/src/yolo/yolo_dataset.py +++ b/api/src/yolo/yolo_dataset.py @@ -66,9 +66,10 @@ class YoloDataset(Dataset): A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.2), A.Affine( - translate_percent=(-0.1, 0.1), - scale=(1.2, 0.8), + translate_percent=(-0.4, 0.4), + scale=(0.5, 1.5), rotate=0, + border_mode=cv2.BORDER_REPLICATE, p=0.5 ), A.Resize(img_size, img_size) @@ -76,7 +77,7 @@ class YoloDataset(Dataset): bbox_params=A.BboxParams( format="coco", label_fields=["labels"], - min_visibility=0.3 + min_visibility=0.5 ) ) else: @@ -102,16 +103,17 @@ class YoloDataset(Dataset): image = cv2.imread(image_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - # Load annotations annotation_ids = self.coco.getAnnIds(imgIds=image_id) annotations = self.coco.loadAnns(annotation_ids) boxes = [] labels = [] + annotations = sorted(annotations, key=lambda x: x['bbox'][2] * x['bbox'][3], reverse=True) + for obj in annotations: xmin, ymin, width, height = obj['bbox'] - xmin, ymin, width, height = float(xmin), float(ymin), float(width), float(height) + xmin, ymin, width, height = int(float(xmin)), int(float(ymin)), int(float(width)), int(float(height)) boxes.append([xmin, ymin, width, height]) labels.append(obj['category_id'] - 1) @@ -126,6 +128,9 @@ class YoloDataset(Dataset): boxes = augmented["bboxes"] labels = augmented["labels"] + if len(boxes) == 0: + return self.__getitem__((idx + 1) % len(self.image_ids)) + ground_truth = list(zip(boxes, labels)) #[S, S, (x+y+w+h+c+C)] targets = np.zeros((self.grid, self.grid, 5 + self.num_classes), dtype=np.float32) diff --git a/api/src/yolo/yolo_loss.py b/api/src/yolo/yolo_loss.py index 2599914..015d9dc 100644 --- a/api/src/yolo/yolo_loss.py +++ b/api/src/yolo/yolo_loss.py @@ -6,7 +6,7 @@ class YoloLoss(nn.Module): def __init__(self): super(YoloLoss, self).__init__() - def forward(self, predictions, targets, lambda_coord=5, lambda_noobj=0.5): + def forward(self, predictions, targets, lambda_coord=1, lambda_noobj=1): pred_boxes = predictions[..., :4] pred_conf = predictions[..., 4] pred_classes = predictions[..., 5:] diff --git a/api/src/yolo/yolo_model.py b/api/src/yolo/yolo_model.py index 0cd2d40..c86ae0f 100644 --- a/api/src/yolo/yolo_model.py +++ b/api/src/yolo/yolo_model.py @@ -1,4 +1,5 @@ import os +import time import torch from tqdm import tqdm import torch.nn as nn @@ -23,25 +24,25 @@ class SkipBlock(nn.Module): 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.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.BatchNorm2d(c_out), - nn.ReLU(inplace=True), + 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.BatchNorm2d(c_out), - nn.ReLU(inplace=True) + 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.BatchNorm2d(c_out), - nn.ReLU(inplace=True) + nn.GroupNorm(num_groups=c_out//8, num_channels=c_out), + nn.LeakyReLU(inplace=True) ) def forward(self, x): - return(F.dropout(F.relu(self.conv_skip(x) + self.conv(x), inplace=True), p=0.3)) + return(F.dropout(F.leaky_relu(self.conv_skip(x) + self.conv(x), inplace=True), p=0.3)) class Yolo_model(nn.Module): @@ -50,8 +51,8 @@ class Yolo_model(nn.Module): self.model = nn.Sequential( nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1), - nn.BatchNorm2d(c_hidden), - nn.ReLU(inplace=True), + 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), @@ -63,25 +64,33 @@ class Yolo_model(nn.Module): 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.BatchNorm2d(c_hidden*4), - nn.ReLU(inplace=True), + 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*4, boxes*5 + labels, kernel_size=1), - nn.Sigmoid() + nn.Conv2d(c_hidden*8, boxes*5 + labels, kernel_size=1) ) def forward(self, x): - return self.model(x).permute(0, 2, 3, 1) + 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) 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(200): + for epoch in range(500): ############ # Training # ############ @@ -137,7 +146,7 @@ def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, mo 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}k | val loss: {int(val_loss * 100000) / 100}k") + 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 @@ -149,5 +158,8 @@ def sample(model, img, device, SAVE_PATH, model_name="face_detection_yolo", fold 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