marvin/CNN_Website
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

working yolo detection model

Browse Source
This commit is contained in:
marvin
2026-05-24 21:38:20 +02:00
parent 87f6bcf7f7
commit 37277319ec
6 changed files with 393 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/src/main.py
+5 -1
View File
@@ -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
+16
View File
@@ -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
+162 -40
View File
@@ -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")
optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, model_name="face_detection_yolo")
api/src/yolo/yolo_dataset.py
+93 -34
View File
@@ -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
api/src/yolo/yolo_loss.py
+33 -4
View File
@@ -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)
noobj_loss = lambda_noobj * torch.mean((pred_conf[noobj_mask]) ** 2)
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)
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
+78 -33
View File
@@ -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))
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