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

detection model works

Browse Source
This commit is contained in:
marvin
2026-05-29 08:32:24 +02:00
parent 37277319ec
commit eec4a7fb10
5 changed files with 60 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/src/util.py
+1 -1
View File
@@ -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):
api/src/yolo/train_yolo_faces.py
+17 -23
View File
@@ -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,
api/src/yolo/yolo_dataset.py
+10 -5
View File
@@ -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)
api/src/yolo/yolo_loss.py
+1 -1
View File
@@ -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:]
api/src/yolo/yolo_model.py
+31 -19
View File
@@ -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