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)