From 87f6bcf7f7fcc65b51c684812488d3a2a544cab3 Mon Sep 17 00:00:00 2001 From: marvinkrausser Date: Fri, 22 May 2026 12:04:50 +0200 Subject: [PATCH] added yolo data viewer --- api/src/yolo/train_yolo_faces.py | 32 ++++++++++++++++++++++++++++++-- api/src/yolo/yolo_dataset.py | 21 +++++++++++---------- 2 files changed, 41 insertions(+), 12 deletions(-) diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py index bf74110..adc4b86 100644 --- a/api/src/yolo/train_yolo_faces.py +++ b/api/src/yolo/train_yolo_faces.py @@ -12,13 +12,38 @@ 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 + +def view_data(dataloader): + 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] + + 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: + 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) def train_yolo(): SAVE_PATH = "./saved_models" IMAGE_SIZE = 64 GRID = 9 - BATCH_SIZE = 256 + BATCH_SIZE = 1 transform = transforms.Compose([ transforms.ToTensor() @@ -37,7 +62,7 @@ def train_yolo(): train_subset, val_subset = random_split(dataset, [train_size, val_size]) - train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=False) + train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True) 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") @@ -48,6 +73,9 @@ def train_yolo(): optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4) loss_module = YoloLoss() + #view_data(train_loader) + #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, diff --git a/api/src/yolo/yolo_dataset.py b/api/src/yolo/yolo_dataset.py index bbf4b53..8eab63d 100644 --- a/api/src/yolo/yolo_dataset.py +++ b/api/src/yolo/yolo_dataset.py @@ -1,3 +1,4 @@ +import math import os import torch from PIL import Image @@ -59,10 +60,10 @@ class YoloDataset(Dataset): 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) + 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]) labels.append(obj['category_id'] - 1) @@ -81,12 +82,12 @@ class YoloDataset(Dataset): 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, 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] del ground_truth[i]