def main():
model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
loss_fn = YoloLoss()
train_dataset = VOCDataset("data/train.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR)
test_dataset = VOCDataset("data/test.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR)
train_loader=DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=1, pin_memory=PIN_MEMORY, shuffle=True,drop_last=True)
test_loader=DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, num_workers=1, pin_memory=PIN_MEMORY, shuffle=True,drop_last=True)
for epoch in range(EPOCHS):
pred_boxes, target_boxes = get_bboxes(train_loader, model, iou_threshold=0.5, threshold=0.4)
mAP = mean_average_precision(pred_boxes, target_boxes, iou_threshold=0.5)
print(f"Train mAP:{mAP}")
train_fn(train_loader, model, optimizer, loss_fn)
if epoch > 99:
for x, y in test_loader:
x = x.to(DEVICE)
for idx in range(16):
bboxes = cellboxes_to_boxes(model(x))
bboxes = non_max_suppression(bboxes[idx], iou_threshold=0.5, threshold=0.4)
plot_image(x[idx].permute(1,2,0).to("cpu"), bboxes)
if __name__ == "__main__":
main()
def train(args):
model = Yolov1(backbone_name=args.backbone)
optimizer = optim.SGD(model.parameters(),
lr=LEARNING_RATE,
momentum=MOMENTUM,
weight_decay=WEIGHT_DECAY)
scheduler = WarmUpMultiStepLR(optimizer,
milestones=STEP_LR_SIZES,
gamma=STEP_LR_GAMMA,
factor=WARM_UP_FACTOR,
num_iters=WARM_UP_NUM_ITERS)
step = model.load_model(optimizer=optimizer, lr_scheduler=scheduler)
model.to(DEVICE)
model.train()
criterion = Yolov1Loss()
while step < NUM_STEPS_TO_FINISH:
data_set = get_voc_data_set(args, percent_coord=True)
t1 = time.perf_counter()
for _, (imgs, gt_boxes, gt_labels, gt_outs) in enumerate(data_set):
step += 1
scheduler.step()
imgs = imgs.to(DEVICE)
gt_outs = gt_outs.to(DEVICE)
model_outs = model(imgs)
loss = criterion(model_outs, gt_outs)
optimizer.zero_grad()
loss.backward()
optimizer.step()
t2 = time.perf_counter()
print('step:{} | loss:{:.8f} | time:{:.4f}'.format(
step, loss.item(), t2 - t1))
t1 = time.perf_counter()
if step != 0 and step % args.save_step == 0:
model.save_model(step, optimizer, scheduler)
def main():
model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE)
optimizer = optim.Adam(
model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY
)
loss_fn = YoloLoss()
if LOAD_MODEL:
load_checkpoint(torch.load(LOAD_MODEL_FILE), model, optimizer)
train_dataset = VOCDataset(
"data/100examples.csv",
transform=transform,
img_dir=IMG_DIR,
label_dir=LABEL_DIR
)
test_dataset = VOCDataset(
"data/test.csv",
transform=transform,
img_dir=IMG_DIR,
label_dir=LABEL_DIR
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True
)
for epoch in range(EPOCHS):
pred_boxes, target_boxes = get_bboxes(
train_loader, model, iou_threshold=0.5, threshold=0.4
)
mean_avg_prec = mean_average_precision(
pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint"
)
print(f"Train mAP in {epoch}: {mean_avg_prec}")
train_fn(train_loader, model, optimizer, loss_fn)
def main():
model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE)
optimizer = optim.Adam(model.parameters(),
lr=LEARNING_RATE,
weight_decay=WEIGHT_DECAY)
loss_fn = YoloLoss()
if LOAD_MODEL:
pass
train_transform = transforms.Compose(
[transforms.Resize(size=(448, 448)),
transforms.ToTensor()])
train_dataset = VOCDataset(csv_file='',
img_root=IMG_DIR,
S=7,
B=2,
C=20,
transform=train_transform)
test_transform = transforms.Compose(
[transforms.Resize(size=(448, 448)),
transforms.ToTensor()])
test_dataset = VOCDataset(csv_file='',
img_root=IMG_DIR,
transform=test_transform)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True,
)
test_loader = DataLoader(dataset=test_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True)
for epoch in range(EPOCHS):
train_fn(train_loader, model, optimizer, loss_fn)
import torch from model import Yolov1 DEVICE = "cuda" if torch.cuda.is_available else "cpu" model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE) print(model) output = model(torch.randn(1, 3, 448, 448).cuda()) print(output.shape) predictions = output.reshape(-1, 7, 7, 20 + 2 * 5) print(predictions.shape)
def main():
model = Yolov1(split_size=S, num_boxes=B, num_classes=C).to(DEVICE)
optimizer = optim.Adam(model.parameters(),
lr=LEARNING_RATE,
weight_decay=WEIGHT_DECAY)
loss_fn = YoloLoss(S=S, B=B, C=C)
if LOAD_MODEL:
load_checkpoint(torch.load(LOAD_MODEL_FILE), model, optimizer)
train_dataset = VOCDataset(
training_path=
'/home/mt/Desktop/For_github/computer_vision_projects/face_recognition/data',
S=3,
C=2,
transform=transform)
# test_dataset = VOCDataset(
# "data/test.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR,
# )
train_loader = DataLoader(
dataset=train_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True,
)
# test_loader = DataLoader(
# dataset=test_dataset,
# batch_size=BATCH_SIZE,
# num_workers=NUM_WORKERS,
# pin_memory=PIN_MEMORY,
# shuffle=True,
# drop_last=True,
# )
for epoch in range(EPOCHS):
# for x, y in train_loader:
# x = x.to(DEVICE)
# for idx in range(8):
# bboxes = cellboxes_to_boxes(model(x))
# bboxes = non_max_suppression(bboxes[idx], iou_threshold=0.5, threshold=0.4, box_format="midpoint")
# plot_image(x[idx].permute(1,2,0).to("cpu"), bboxes)
# import sys
# sys.exit()
pred_boxes, target_boxes = get_bboxes(train_loader,
model,
iou_threshold=0.5,
threshold=0.4)
mean_avg_prec = mean_average_precision(pred_boxes,
target_boxes,
iou_threshold=0.5,
box_format="midpoint")
print(f"Train mAP: {mean_avg_prec}")
#if mean_avg_prec > 0.9:
# checkpoint = {
# "state_dict": model.state_dict(),
# "optimizer": optimizer.state_dict(),
# }
# save_checkpoint(checkpoint, filename=LOAD_MODEL_FILE)
# import time
# time.sleep(10)
train_fn(train_loader, model, optimizer, loss_fn)
def main_():
model = Yolov1(split_size=Grid, num_boxes=BBOX,
num_classes=C_num).to(DEVICE)
model_opt = Yolov1(split_size=Grid, num_boxes=BBOX,
num_classes=C_num).to(DEVICE)
best_loss = 1e6
# print(model)
optimizer = optim.Adam(model.parameters(),
lr=LEARNING_RATE,
weight_decay=WEIGHT_DECAY)
loss_fn = YoloLoss()
train_dataset = YoloDataSet(
"data/train.csv",
transform=transform,
img_dir=IMG_DIR,
label_dir=LABEL_DIR,
)
test_dataset = YoloDataSet(
"data/test.csv",
transform=transform,
img_dir=IMG_DIR,
label_dir=LABEL_DIR,
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True,
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True,
)
test_class = []
test_conf = []
test_win = []
test_noob = []
test_cen = []
train_class = []
train_conf = []
train_win = []
train_noob = []
train_cen = []
for epoch in range(EPOCHS):
print("epoch: " + str(epoch))
mean_loss = []
mean_coloss = []
mean_clloss = []
mean_celoss = []
mean_wloss = []
mean_noloss = []
model.eval()
for batch_idx, (x, y) in enumerate(test_loader):
x, y = x.to(DEVICE), y.to(DEVICE)
with torch.no_grad():
out = model(x)
# x, y = x_train.to(DEVICE), y_train.to(DEVICE)
# out = model(x)
loss, confidence_loss, class_loss, center_loss, window_loss, noobj_loss = loss_fn(
out, y)
mean_loss.append(loss.item())
mean_coloss.append(confidence_loss.item())
mean_clloss.append(class_loss.item())
mean_celoss.append(center_loss.item())
mean_wloss.append(window_loss.item())
mean_noloss.append(noobj_loss.item())
test_loss = sum(mean_loss) / len(mean_loss)
if test_loss < best_loss:
model_opt = model
best_loss = test_loss
model.train()
print(f"test loss was {test_loss}")
# print(f"Mean confidance loss was {sum(mean_coloss)/len(mean_coloss)}")
# print(f"Mean class loss was {sum(mean_clloss)/len(mean_clloss)}")
# print(f"Mean center loss was {sum(mean_celoss)/len(mean_celoss)}")
# print(f"Mean window loss was {sum(mean_wloss)/len(mean_wloss)}")
# print(f"Mean noobj loss was {sum(mean_noloss)/len(mean_noloss)}")
test_class.append(sum(mean_clloss) / len(mean_clloss))
test_conf.append(sum(mean_coloss) / len(mean_coloss))
test_cen.append(sum(mean_celoss) / len(mean_celoss))
test_win.append(sum(mean_wloss) / len(mean_wloss))
test_noob.append(sum(mean_noloss) / len(mean_noloss))
# pred_boxes, target_boxes = get_bboxes(test_loader, model, iou_threshold=0.5, threshold=0.4, S=Grid, B=BBOX, C=C_num)
# mean_avg_prec = mean_average_precision_built(pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint")
# print(f"test mAP: {mean_avg_prec}")
# pred_boxes, target_boxes = get_bboxes(train_loader, model, iou_threshold=0.5, threshold=0.4, S=Grid, B=BBOX, C=C_num)
# mean_avg_prec = mean_average_precision_built(pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint")
# print(f"Train mAP: {mean_avg_prec}")
conf_loss, cla_loss, cen_loss, win_loss, noo_los = train_fn(
train_loader, model, optimizer, loss_fn)
train_class.append(cla_loss)
train_conf.append(conf_loss)
train_cen.append(cen_loss)
train_win.append(win_loss)
train_noob.append(noo_los)
ploLoss(epoch, train_class, test_class, "class")
ploLoss(epoch, train_conf, test_conf, "confidence")
ploLoss(epoch, train_win, test_win, "win")
ploLoss(epoch, train_noob, test_noob, "noobj")
ploLoss(epoch, train_cen, test_cen, "center")
model.eval()
torch.save(model_opt, PATH)
def main():
model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE)
optimizer = optim.Adam(
model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY
)
loss_fn = YoloLoss()
if LOAD_MODEL:
load_checkpoint(torch.load(LOAD_MODEL_FILE), model, optimizer)
train_dataset = VOCDataset(
# test.csv, 8examples.csv, 100examples.csv
"data/8examples.csv",
transform=transform,
img_dir=IMG_DIR,
label_dir=LABEL_DIR,
)
test_dataset = VOCDataset(
"data/test.csv" , transform = transform, img_dir = IMG_DIR, label_dir = LABEL_DIR
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=False,
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=False,
)
for epoch in range(EPOCHS):
# for x, y in train_loader:
# x = x.to(DEVICE)
# for idx in range(8):
# bboxes = cellboxes_to_boxes(model(x))
# bboxes = non_max_suppression(bboxes[idx], iou_threshold=0.5, threshold=0.4, box_format="midpoint")
# plot_image(x[idx].permute(1,2,0).to("cpu"), bboxes)
# import sys
# sys.exit()
pred_boxes, target_boxes = get_bboxes(
train_loader, model, iou_threshold=0.5, threshold=0.4 , device= DEVICE,
)
mean_avg_prec = mean_average_precision(
pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint"
)
if mean_avg_prec > 0.9:
checkpoint = {
"state_dict": model.state_dict(),
"optimizer": optimizer.state_dict(),
}
save_checkpoint(checkpoint, filename=LOAD_MODEL_FILE)
import time
time.sleep(10)
print(f"Train mAP: {mean_avg_prec}")
train_fn(train_loader, model, optimizer, loss_fn)
import torch
from config import GRID_NUM, DEVICE
from dataset import VOC_CLASSES
from model import Yolov1
from matplotlib import pyplot as plt
from torchvision import transforms
import PIL
from PIL import Image
import numpy as np
from os import path as osp
import os
from PIL import Image
model = Yolov1(backbone_name='resnet101')
model.load_model()
def draw_box(img_np,
boxes_np,
tags_np,
scores_np=None,
relative_coord=False,
save_path=None):
if scores_np is None:
scores_np = [1.0 for i in tags_np]
# img = cv2.cvtColor(img_cv, cv2.COLOR_RGB2GRAY)
# print(f'img_np.shape:{img_np.shape}')
h, w, _ = img_np.shape
if relative_coord:
boxes_np = np.array([
boxes_np[:, 0] * w,
epoch_cls_loss += cls_loss.data[0] / len(test_loader)
return epoch_loss, epoch_loc_loss, epoch_cnf_loss, epoch_cls_loss
if __name__ == '__main__':
parse = argparse.ArgumentParser()
parse.add_argument("-d", "--debug", action="store_true")
parse.add_argument("-r", "--res", action="store_false")
parse.add_argument("-e", "--epoch", type=int, default=1000)
arg = parse.parse_args()
VOC_BASE = os.path.join("..", "data", "VOC2012")
IMAGE_SIZE = 300
BATCH_SIZE = 64
model = Yolov1(IMAGE_SIZE) if arg.res else YOLORes(IMAGE_SIZE)
GRID_NUM = model.grid_size
viz = visdom.Visdom(port=6006)
if cuda:
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
train_a, test_a = get_annotations(VOC_BASE)
train_loader = DataLoader(VocDataSet(train_a,
image_size=IMAGE_SIZE,
grid_num=GRID_NUM,
dir=VOC_BASE),
batch_size=BATCH_SIZE,
num_workers=4)
scores = torch.stack(scores, 0).numpy() # .squeeze(dim=0)
# print(f'*** boxes:{boxes}\ntags:{tags}\nscores:{scores}')
else:
boxes = torch.FloatTensor([])
tags = torch.LongTensor([]) # .squeeze(dim=0)
scores = torch.FloatTensor([]) # .squeeze(dim=0)
# img, boxes, tags, scores = np.array(img), np.array(boxes), np.array(tags), np.array(scores)
return boxes, tags, scores
if __name__ == '__main__':
# test:
# fake_pred = torch.rand(1, GRID_NUM, GRID_NUM, 30)
# decoder(fake_pred)
CONTINUE = False # continue from breakpoint
model = Yolov1(backbone_name='resnet50')
model.load_model()
# predict_one_img('../test_img/000001.jpg', model)
# test_img_dir = '../test_img'
test_img_dir = '/Users/chenlinwei/Dataset/VOC0712/VOC2012test/JPEGImages'
for root, dirs, files in os.walk(test_img_dir, topdown=True):
if test_img_dir == root:
print(root, dirs, files)
files = [
i for i in files
if any([j in i for j in ['jpg', 'png', 'jpeg', 'gif', 'tiff']])
]
shuffle(files)
if CONTINUE:
with open(osp.join(test_img_dir, 'tested.txt'), 'a') as _:
pass
def setUp(self):
self.model = Yolov1(grid_size=7, num_boxes=2, num_classes=20)
