diff --git a/api/src/main.py b/api/src/main.py index ea02671..4917b26 100644 --- a/api/src/main.py +++ b/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() \ No newline at end of file +def main(): + train_yolo() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/api/src/util.py b/api/src/util.py index 7e709ea..127ed1b 100644 --- a/api/src/util.py +++ b/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): @@ -55,4 +59,16 @@ def iou(boxA, boxB): union = boxA_area + boxB_area - inter_area - return inter_area / union if union > 0 else 0 \ No newline at end of file + 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)) \ 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 adc4b86..1528360 100644 --- a/api/src/yolo/train_yolo_faces.py +++ b/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_model import train, Yolo_model +from yolo.yolo_dataset import YoloDataset, turn_image_centered +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] - grid_size = image_size // label.shape[0] + visualize_boxes(label, image) - boxes_to_draw = [] - grids_to_draw = [] - for x in range(label.shape[0]): - for y in range(label.shape[1]): - grids_to_draw.append([x*grid_size, y*grid_size, (x+1)*grid_size, (y+1)*grid_size]) - if label[x, y, 4].item() == 0: - continue - boxes_to_draw.append(label[x, y, :4]) - if len(boxes_to_draw) == 0: +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) + 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 = torch.tensor(grids_to_draw) - image = draw_bounding_boxes(image, grids_to_draw, colors=(0, 255, 0)) - image = draw_bounding_boxes(image, boxes_to_draw, colors=(255, 0, 0)) - visualizeImage(image) + 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(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 - BATCH_SIZE = 1 - - transform = transforms.Compose([ - transforms.ToTensor() - ]) + GRID = 3 + BATCH_SIZE = 32 dataset = YoloDataset( - image_dir="data/faces/train", - annotation_path="data/faces/train/_annotations.coco.json", + image_dir="data/faces_2/train", + 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)) - val_size = len(dataset) - train_size - - train_subset, val_subset = random_split(dataset, [train_size, val_size]) + dataset_valid = YoloDataset( + image_dir="data/faces_2/test", + annotation_path="data/faces_2/test/_annotations.coco.json", + img_size=IMAGE_SIZE, + transform=False, + grid=GRID + ) - train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True) - val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False) + train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0) + 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, - model_name="face_detection_yolo") \ No newline at end of file + 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_dataset.py b/api/src/yolo/yolo_dataset.py index 8eab63d..d824def 100644 --- a/api/src/yolo/yolo_dataset.py +++ b/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_j = int(x / cell_w) + 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) @@ -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") - - orig_w, orig_h = image.size - - scale_w = self.img_size / orig_w - scale_h = self.img_size / orig_h + image = cv2.imread(image_path) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 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) - 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]) + boxes.append([xmin, ymin, width, height]) 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) - targets[x, y, 0] = box[0] #x - targets[x, y, 1] = box[1] #y - targets[x, y, 2] = box[2] #w - targets[x, y, 3] = box[3] #h - targets[x, y, 4] = 1 #confidence - for i in range(len(class_one_hot)): #label - targets[x, y, i + 5] = class_one_hot[i] + 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, + img_h=self.img_size, S=self.grid, cell_i=x, cell_j=y) + + + targets[x, y, 0] = x_grid_centered / (self.img_size / self.grid) + targets[x, y, 1] = y_grid_centered / (self.img_size / self.grid) + + 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) - - # resize image - image = image.resize((self.img_size, self.img_size)) - - if self.transforms: - image = self.transforms(image) + image = self.toTensor(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 d27a572..2599914 100644 --- a/api/src/yolo/yolo_loss.py +++ b/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): @@ -15,13 +16,41 @@ class YoloLoss(nn.Module): obj_mask = targets[..., 4] == 1 noobj_mask = targets[..., 4] == 0 + + + if obj_mask.any(): + 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) + - 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) - noobj_loss = lambda_noobj * torch.mean((pred_conf[noobj_mask]) ** 2) - - 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 + + return total_loss \ No newline at end of file diff --git a/api/src/yolo/yolo_model.py b/api/src/yolo/yolo_model.py index 76bf15b..0cd2d40 100644 --- a/api/src/yolo/yolo_model.py +++ b/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_model(nn.Module): - def __init__(self, c_in, c_hidden, boxes, img_size, grid, labels): +class Yolo_Conv_Block(nn.Module): + def __init__(self, c_in, c_hidden, c_out, kernel_size): 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.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) + nn.Conv2d(in_channels=c_hidden, out_channels=c_out, kernel_size=1) ) def forward(self, x): - batch = x.shape[0] - x = self.model(x) - x = x.reshape(batch, self.grid, self.grid, -1) - return 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.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)) + + +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, device="cpu"): +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 \ No newline at end of file + 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 \ No newline at end of file