working yolo detection model

api/src/main.py

@@ -2,4 +2,8 @@ from r_cnn.r_cnn_test import train_cnn_test
 from yolo.train_yolo_faces import train_yolo
 
 
-train_yolo()
+def main():
+    train_yolo()
+
+if __name__ == "__main__":
+    main()

api/src/util.py

@@ -1,5 +1,9 @@
 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):
@@ -56,3 +60,15 @@ def iou(boxA, boxB):
     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=16,pin_memory=True)
+        start = time.time()
+        for _ in range(2):
+            for images, _  in tqdm(train_loader, leave=False):
+                images = images.to(device)
+                _ = model(images)
+        end = time.time()
+        print("Finish with:{} seconds, num_workers={}".format(int(end - start), num_workers))

api/src/yolo/train_yolo_faces.py

@@ -1,11 +1,11 @@
 import os
 
-from torchvision import transforms
+import cv2
 from tqdm import tqdm
-from yolo.yolo_dataset import YoloDataset
+from yolo.yolo_dataset import YoloDataset, turn_image_centered
-from yolo.yolo_model import train, Yolo_model
+from yolo.yolo_model import train, Yolo_model, sample
 from yolo.yolo_loss import YoloLoss
-from util import TransformedSubset, visualizeImage
+from util import TransformedSubset, test_workers_speed, visualizeImage
 from torch.utils.data import DataLoader, random_split
 from torchvision import datasets
 import torch
@@ -13,70 +13,192 @@ 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 view_data(dataloader):
+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 view_data(dataset):
+    dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=False)
     for images, labels in iter(dataloader):
         for batch in range(images.shape[0]):
             image = images[batch]
             label = labels[batch]
-            image_size = image.shape[1]
+            visualize_boxes(label, image)
-            grid_size = image_size // label.shape[0]
 
-            boxes_to_draw = []
+def sample_data(dataloader, model, device, SAVE_PATH):
-            grids_to_draw = []
+    for images, labels in iter(dataloader):
-            for x in range(label.shape[0]):
+        start = time.perf_counter()
-                for y in range(label.shape[1]):
+        predictions = sample(model, images, device, SAVE_PATH)
-                    grids_to_draw.append([x*grid_size, y*grid_size, (x+1)*grid_size, (y+1)*grid_size])
+        end = time.perf_counter()
-                    if label[x, y, 4].item() == 0:
+        print(f"Elapsed: {end - start:.6f} seconds") 
-                        continue
+        for batch in range(predictions.shape[0]):
-                    boxes_to_draw.append(label[x, y, :4])
+            image = images[batch]
-            if len(boxes_to_draw) == 0:
+            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)
+    if len(boxes_to_draw) == 0:
+        print("no labels")
+        return
+    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 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
-            boxes_to_draw = torch.stack(boxes_to_draw)
+            grids_to_draw_obj.append([x*grid_size, y*grid_size, (x+1)*grid_size, (y+1)*grid_size]) #xmin, ymin, xmax, ymax
-            grids_to_draw = torch.tensor(grids_to_draw)
+
-            image = draw_bounding_boxes(image, grids_to_draw, colors=(0, 255, 0))
+            boxx = label[x, y, 0] * (image_size / grid_number)
-            image = draw_bounding_boxes(image, boxes_to_draw, colors=(255, 0, 0))
+            boxy = label[x, y, 1] * (image_size / grid_number)
-            visualizeImage(image)
+
+            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(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():
+    # 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)
+    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 / 64
+        scale_h = H / 64
+
+        image = image / 255.0
+        
+        image = image.unsqueeze(0)
+        image = F.interpolate(image, size=(64, 64), mode="bilinear", align_corners=False) #different modes?
+
+        prediction = model(image)
+
+        bboxes, grid_ob, grid_noob = convert_prediction(prediction.squeeze(0), image.squeeze(0), threshold=0.995)
+
+        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 = 9
+    GRID = 3
-    BATCH_SIZE = 1
+    BATCH_SIZE = 32
-
-    transform = transforms.Compose([
-        transforms.ToTensor()
-    ])
 
     dataset = YoloDataset(
-        image_dir="data/faces/train", 
+        image_dir="data/faces_2/train", 
-        annotation_path="data/faces/train/_annotations.coco.json",
+        annotation_path="data/faces_2/train/_annotations.coco.json",
         img_size=IMAGE_SIZE,
-        transforms=transform
+        transform=True,
+        grid=GRID
     )
 
-    train_size = int(0.8 * len(dataset))
+    dataset_valid = YoloDataset(
-    val_size = len(dataset) - train_size
+        image_dir="data/faces_2/test", 
-
+        annotation_path="data/faces_2/test/_annotations.coco.json",
-    train_subset, val_subset = random_split(dataset, [train_size, val_size])
+        img_size=IMAGE_SIZE,
+        transform=False,
+        grid=GRID
+    )
 
 
-    train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True)
+    train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
-    val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False)
+    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, c_hidden=32, boxes=1, img_size=IMAGE_SIZE, grid=GRID, labels=1)
+    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(train_loader)
+
+    #view_data(dataset)
+    #exit()
+
+
+    sample_data(dataloader=val_loader, model=model, device=device, SAVE_PATH=SAVE_PATH)
+    exit()
+
+
+    #test_workers_speed(dataset, model)
+    #exit()
+
+    #use_webcam()
     #exit()
 
 
     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, 
+                optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, model_name="face_detection_yolo")
-                model_name="face_detection_yolo")

api/src/yolo/yolo_dataset.py

@@ -1,10 +1,13 @@
 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
 
 
 def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
@@ -12,8 +15,8 @@ def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
         cell_h = img_h / S
 
         # find which cell the center belongs to
-        gt_cell_i = int(y / cell_h)
+        gt_cell_i = min(int(x / cell_w), S - 1)
-        gt_cell_j = int(x / cell_w)
+        gt_cell_j = min(int(y / cell_h), S - 1)
 
         return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
 
@@ -25,14 +28,68 @@ 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 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 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, transforms=None):
+    def __init__(self, image_dir, annotation_path, img_size=64, grid = 9, transform=False):
         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
+        self.num_classes = len(self.coco.cats)-1
+        self.toTensor = transforms.ToTensor()
+
+        if transform:
+            self.transform = A.Compose(
+                    [
+                        A.HorizontalFlip(p=0.5),
+                        A.RandomBrightnessContrast(p=0.2),
+                        A.Affine(
+                            translate_percent=(-0.1, 0.1),
+                            scale=(1.2, 0.8),
+                            rotate=0,
+                            p=0.5
+                        ),
+                        A.Resize(img_size, img_size)
+                    ],
+                    bbox_params=A.BboxParams(
+                        format="coco",
+                        label_fields=["labels"],
+                        min_visibility=0.3
+                )
+            )
+        else:
+            self.transform = A.Compose(
+                    [
+                        A.Resize(img_size, img_size)
+                    ],
+                    bbox_params=A.BboxParams(
+                        format="coco",
+                        label_fields=["labels"],
+                        min_visibility=0.3
+                )
+            )
 
     def __len__(self):
         return len(self.image_ids)
@@ -42,12 +99,8 @@ class YoloDataset(Dataset):
         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")
+        image =  cv2.imread(image_path)
-
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        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)
@@ -60,47 +113,53 @@ class YoloDataset(Dataset):
             xmin, ymin, width, height = obj['bbox']
             xmin, ymin, width, height = float(xmin), float(ymin), float(width), float(height)
 
-            xmin = math.ceil(xmin * scale_w)
+            boxes.append([xmin, ymin, width, height])
-            ymin = math.ceil(ymin * scale_h)
-            xmax = math.floor(xmin + max(width * scale_w, 1))
-            ymax = math.floor(ymin + max(height * scale_h, 1))
-
-            boxes.append([xmin, ymin, xmax, ymax])
             labels.append(obj['category_id'] - 1)
 
+        augmented = self.transform(
+            image=image,
+            bboxes=boxes,
+            labels=labels
+        )
+
+        image = augmented["image"]
+        boxes = augmented["bboxes"]
+        labels = augmented["labels"]
+
         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))
+        targets = np.zeros((self.grid, self.grid, 5 + self.num_classes), dtype=np.float32)
         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, 
+                    x_center = box[0] + box[2] / 2
+                    y_center = box[1] + box[3] / 2
+
+                    if is_center_in_grid_cell(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):
                         
-                        class_one_hot = one_hot(label, num_labels)
+                        class_one_hot = one_hot(int(label), self.num_classes)
-                        targets[x, y, 0] = box[0]  #x
+                        x_grid_centered, y_grid_centered = turn_grid_centered(x=x_center, y=y_center, img_w=self.img_size, 
-                        targets[x, y, 1] = box[1]  #y
+                                              img_h=self.img_size, S=self.grid, cell_i=x, cell_j=y)
-                        targets[x, y, 2] = box[2]  #w
+
-                        targets[x, y, 3] = box[3]  #h
+
-                        targets[x, y, 4] = 1       #confidence
+                        targets[x, y, 0] = x_grid_centered / (self.img_size / self.grid)
-                        for i in range(len(class_one_hot)):  #label
+                        targets[x, y, 1] = y_grid_centered / (self.img_size / self.grid)
-                            targets[x, y, i + 5] = class_one_hot[i]
+
+                        targets[x, y, 2] = box[2] / self.img_size
+                        targets[x, y, 3] = box[3] / self.img_size
+                        targets[x, y, 4] = 1
+                        for j in range(len(class_one_hot)):
+                            targets[x, y, j + 5] = class_one_hot[j]
 
                         del ground_truth[i]
-
                         break
 
 
         targets = torch.tensor(targets, dtype=torch.float32)
-
+        image = self.toTensor(image)
-        # resize image
-        image = image.resize((self.img_size, self.img_size))
-
-        if self.transforms:
-            image = self.transforms(image)
 
         return image, targets

api/src/yolo/yolo_loss.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 class YoloLoss(nn.Module):
     def __init__(self):
@@ -16,12 +17,40 @@ class YoloLoss(nn.Module):
         obj_mask = targets[..., 4] == 1
         noobj_mask = targets[..., 4] == 0
 
-        box_loss = lambda_coord * torch.mean((pred_boxes[obj_mask] - target_boxes[obj_mask]) ** 2)
 
-        obj_loss = torch.mean((pred_conf[obj_mask] - target_conf[obj_mask]) ** 2)
+        if obj_mask.any():
-        noobj_loss = lambda_noobj * torch.mean((pred_conf[noobj_mask]) ** 2)
+            box_loss = F.mse_loss(
+                pred_boxes[obj_mask],
+                target_boxes[obj_mask],
+                reduction="mean"
+            )
+        else:
+            box_loss = torch.tensor(0.0, device=predictions.device)
+
+        box_loss = lambda_coord * box_loss
+
+
+        obj_loss = F.mse_loss(
+            pred_conf[obj_mask],
+            target_conf[obj_mask],
+            reduction="mean"
+        ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)
+
+        noobj_loss = F.mse_loss(
+            pred_conf[noobj_mask],
+            target_conf[noobj_mask],
+            reduction="mean"
+        ) if noobj_mask.any() else torch.tensor(0.0, device=predictions.device)
+
+        noobj_loss = lambda_noobj * noobj_loss
+
+
+        class_loss = F.binary_cross_entropy_with_logits(
+            pred_classes[obj_mask],
+            target_classes[obj_mask]
+        ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)
 
-        class_loss = torch.mean((pred_classes[obj_mask] - target_classes[obj_mask]) ** 2)
         
         total_loss = box_loss + obj_loss + noobj_loss + class_loss
+
         return total_loss

api/src/yolo/yolo_model.py

@@ -2,47 +2,84 @@ import os
 import torch
 from tqdm import tqdm
 import torch.nn as nn
+import torch.nn.functional as F
 
-
+class Yolo_Conv_Block(nn.Module):
-class Yolo_model(nn.Module):
+    def __init__(self, c_in, c_hidden, c_out, kernel_size):
-    def __init__(self, c_in, c_hidden, boxes, img_size, grid, labels):
         super().__init__()
 
-        self.grid = grid
-
         self.model = nn.Sequential(
-            nn.Conv2d(in_channels=c_in, out_channels=c_hidden, kernel_size=7, padding=3),
+            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_hidden, kernel_size=3, padding=1),
+            nn.Conv2d(in_channels=c_hidden, out_channels=c_out, kernel_size=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)),
-            nn.Dropout(0.3)
         )
 
     def forward(self, x):
-        batch = x.shape[0]
+        return self.model(x)
-        x = self.model(x)
+    
-        x = x.reshape(batch, self.grid, self.grid, -1)
+class SkipBlock(nn.Module):
-        return x
+    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.dropout(F.relu(self.conv_skip(x) + self.conv(x), inplace=True), p=0.3))
 
 
-def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True, device="cpu"):
+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.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.BatchNorm2d(c_hidden*4),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.3),
+
+            nn.AdaptiveAvgPool2d((grid, grid)),
+            nn.Conv2d(c_hidden*4, boxes*5 + labels, kernel_size=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        return self.model(x).permute(0, 2, 3, 1)
+    
+
+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):
         ############
@@ -88,9 +125,6 @@ def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, mo
                 loss = loss_module(prediction, labels)
                 lossCount += loss.item()
 
-
-                lossCount += loss.sum().item()
-
                 count += images.size(0)
 
         val_loss = lossCount / count
@@ -103,6 +137,17 @@ 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 * 1000) / 100} | val loss: {int(val_loss * 1000) / 100}")
+        print(f"epoch: {epoch+1} | train loss: {int(train_loss * 100000) / 100}k | val loss: {int(val_loss * 100000) / 100}k")
         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)
+        pred = model(img)
+        return pred