changed api structure and added object detection
This commit is contained in:
@@ -0,0 +1,55 @@
|
||||
import os
|
||||
|
||||
from torchvision import transforms
|
||||
from tqdm import tqdm
|
||||
from yolo_model import train, Yolo_model
|
||||
from yolo_loss import YoloLoss
|
||||
from cocoDetectionDataset import CocoDetectionDatasetResized
|
||||
from util import TransformedSubset, 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
|
||||
|
||||
|
||||
|
||||
SAVE_PATH = "./saved_models"
|
||||
IMAGE_SIZE = 128
|
||||
GRID = 9
|
||||
BATCH_SIZE = 1
|
||||
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
dataset = CocoDetectionDatasetResized(
|
||||
image_dir="data/faces/train",
|
||||
annotation_path="data/faces/train/_annotations.coco.json",
|
||||
img_size=IMAGE_SIZE,
|
||||
transforms=transform
|
||||
)
|
||||
|
||||
train_size = int(0.8 * len(dataset))
|
||||
val_size = len(dataset) - train_size
|
||||
|
||||
train_subset, val_subset = random_split(dataset, [train_size, val_size])
|
||||
|
||||
|
||||
train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
|
||||
val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=lambda x: tuple(zip(*x)))
|
||||
|
||||
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=2, img_size=IMAGE_SIZE, grid=GRID, labels=1)
|
||||
model.to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
||||
loss_module = YoloLoss()
|
||||
|
||||
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", img_size=IMAGE_SIZE, num_classes=1, grid=GRID)
|
||||
|
||||
exit()
|
||||
@@ -0,0 +1,107 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class YoloLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super(YoloLoss, self).__init__()
|
||||
|
||||
@staticmethod
|
||||
def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
|
||||
cell_w = img_w / S
|
||||
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)
|
||||
|
||||
return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
|
||||
|
||||
@staticmethod
|
||||
def save_sqrt(i):
|
||||
return torch.sqrt(torch.clamp(i, min=1e-6))
|
||||
|
||||
@staticmethod
|
||||
def iou(boxA, boxB):
|
||||
xA = torch.max(boxA[0], boxB[0])
|
||||
yA = torch.max(boxA[1], boxB[1])
|
||||
xB = torch.min(boxA[2], boxB[2])
|
||||
yB = torch.min(boxA[3], boxB[3])
|
||||
|
||||
inter_area = torch.clamp(xB - xA, min=0) * torch.clamp(yB - yA, min=0)
|
||||
|
||||
boxA_area = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
|
||||
boxB_area = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
|
||||
|
||||
union = boxA_area + boxB_area - inter_area
|
||||
|
||||
return inter_area / (union + 1e-6)
|
||||
|
||||
@staticmethod
|
||||
def xywh_to_xyxy(box):
|
||||
x, y, w, h = box
|
||||
|
||||
return torch.stack([
|
||||
x - w / 2,
|
||||
y - h / 2,
|
||||
x + w / 2,
|
||||
y + h / 2
|
||||
])
|
||||
|
||||
def forward(self, pred, target, labels, cell_size, img_w, img_h, factor_no_object = 0.5, factor_object = 5):
|
||||
#pred: [SxSx(B(x,y,w,h,con)C)]
|
||||
#target: [B(x, y, w, h, C)]
|
||||
|
||||
for s0 in range(pred.shape[0]):
|
||||
for s1 in range(pred.shape[1]):
|
||||
boxes = pred[s0, s1, :(-1*labels)]
|
||||
pred_label = pred[s0, s1, -1:]
|
||||
|
||||
ground_truth_box = None
|
||||
for b in range(target.shape[0] // (4 + labels)):
|
||||
box = target[(b)*(4 + labels):(b+1)*(4 + labels)]
|
||||
if self.is_center_in_grid_cell(box[0]+box[2]/2, box[1]+box[3]/2, img_w, img_h, cell_size, s0, s1):
|
||||
ground_truth_box = box[:4]
|
||||
ground_truth_label = box[4:]
|
||||
break
|
||||
|
||||
|
||||
coord_loss = torch.tensor(0., device=pred.device)
|
||||
size_loss = torch.tensor(0., device=pred.device)
|
||||
prob_loss = torch.tensor(0., device=pred.device)
|
||||
label_loss = torch.tensor(0., device=pred.device)
|
||||
|
||||
if ground_truth_box is not None:
|
||||
|
||||
iou_stats = []
|
||||
|
||||
for b in range(boxes.shape[0] // 5):
|
||||
box = boxes[(b)*5:(b+1)*5]
|
||||
|
||||
iou_stats.append(self.iou(self.xywh_to_xyxy(box[0:4]), self.xywh_to_xyxy(ground_truth_box[0:4])))
|
||||
iou_tensor = torch.stack(iou_stats)
|
||||
responsible = torch.argmax(iou_tensor)
|
||||
|
||||
for b in range(boxes.shape[0] // 5):
|
||||
box = boxes[(b)*5:(b+1)*5]
|
||||
|
||||
is_responsible = (b == responsible).float()
|
||||
|
||||
coord_loss += factor_object * is_responsible * ((box[0] - ground_truth_box[0])**2 + (box[1] - ground_truth_box[1])**2)
|
||||
size_loss += factor_object * is_responsible * ((self.save_sqrt(box[2]) - self.save_sqrt(ground_truth_box[2]))**2 +
|
||||
(self.save_sqrt(box[3]) - self.save_sqrt(ground_truth_box[3]))**2)
|
||||
|
||||
prob_loss += is_responsible * (box[4] - 1)**2
|
||||
prob_loss += factor_no_object * (1 - is_responsible) * box[4]**2
|
||||
|
||||
|
||||
for i in range(pred_label.shape[0]):
|
||||
label_loss += ((pred_label[i] - ground_truth_label[i])**2) / pred_label.shape[0]
|
||||
|
||||
|
||||
else:
|
||||
for b in range(boxes.shape[0] // 5):
|
||||
box = boxes[(b)*5:(b+1)*5]
|
||||
|
||||
prob_loss += factor_no_object * box[4] ** 2
|
||||
|
||||
return (label_loss + prob_loss + coord_loss + size_loss) / pred.shape[0]**2
|
||||
@@ -0,0 +1,134 @@
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms.functional as TF
|
||||
import time
|
||||
|
||||
|
||||
class Yolo_model(nn.Module):
|
||||
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.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Flatten(),
|
||||
nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels))
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
batch = x.shape[0]
|
||||
x = self.model(x)
|
||||
x = x.reshape(batch, self.grid, self.grid, -1)
|
||||
return x
|
||||
|
||||
def create_stack(images, annotations, num_classes, device):
|
||||
data = []
|
||||
labels = []
|
||||
for i in range(len(images)):
|
||||
image = images[i]
|
||||
|
||||
annotation = annotations[i]
|
||||
|
||||
label_train = []
|
||||
for box, label in tuple(zip(annotation["boxes"], annotation["labels"])):
|
||||
box = box
|
||||
label = 0
|
||||
one_hot = F.one_hot(torch.tensor(label, dtype=torch.long), num_classes=num_classes)
|
||||
label_train.append(torch.cat((box, one_hot), dim=0))
|
||||
|
||||
label_torch = torch.stack(label_train).reshape(-1).to(device)
|
||||
labels.append(label_torch)
|
||||
data.append(image)
|
||||
|
||||
return data, labels
|
||||
|
||||
|
||||
def train(model, num_classes, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, grid, saving=True, device="cpu", img_size=64):
|
||||
best_val = torch.finfo(torch.float32).max
|
||||
|
||||
for epoch in range(200):
|
||||
############
|
||||
# Training #
|
||||
############
|
||||
model.train()
|
||||
|
||||
count, lossCount = 0, 0.
|
||||
for images, annotations in tqdm(train_loader, desc=f"Train", leave=False):
|
||||
data, labels = create_stack(images, annotations, num_classes, device)
|
||||
|
||||
data = torch.stack(data).to(device)
|
||||
|
||||
prediction = model(data)
|
||||
|
||||
loss = []
|
||||
for pred, label in tuple(zip(prediction, labels)):
|
||||
loss.append(loss_module(pred = pred, target= label, labels=num_classes, cell_size=grid, img_w = img_size, img_h = img_size))
|
||||
loss = torch.stack(loss).to(device)
|
||||
loss = loss.sum()
|
||||
lossCount += loss.item()
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
|
||||
count += data.size(0)
|
||||
|
||||
train_loss = lossCount / count
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
##############
|
||||
# Validation #
|
||||
##############
|
||||
model.eval()
|
||||
|
||||
count, lossCount = 0, 0.
|
||||
for images, annotations in tqdm(val_loader, desc=f"Test", leave=False):
|
||||
with torch.no_grad():
|
||||
data, labels = create_stack(images, annotations, num_classes, device)
|
||||
|
||||
data = torch.stack(data).to(device)
|
||||
|
||||
prediction = model(data)
|
||||
|
||||
loss = []
|
||||
for pred, label in tuple(zip(prediction, labels)):
|
||||
loss.append(loss_module(pred = pred, target= label, labels=num_classes, cell_size=grid, img_w = img_size, img_h = img_size))
|
||||
loss = torch.stack(loss).to(device)
|
||||
|
||||
|
||||
lossCount += loss.sum().item()
|
||||
|
||||
count += data.size(0)
|
||||
|
||||
val_loss = lossCount / count
|
||||
|
||||
if(saving and best_val > val_loss):
|
||||
best_val = val_loss
|
||||
save_dir = os.path.join(SAVE_PATH, model_name)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
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}")
|
||||
torch.cuda.empty_cache()
|
||||
return best_val
|
||||
Reference in New Issue
Block a user