def main(): """ YOLOv3 trainer. See README for details. """ args = parse_args() print("Setting Arguments.. : ", args) cuda = torch.cuda.is_available() and args.use_cuda os.makedirs(args.checkpoint_dir, exist_ok=True) # Parse config settings with open(args.cfg, 'r') as f: cfg = yaml.load(f) print("successfully loaded config file: ", cfg) momentum = cfg['TRAIN']['MOMENTUM'] decay = cfg['TRAIN']['DECAY'] burn_in = cfg['TRAIN']['BURN_IN'] iter_size = cfg['TRAIN']['MAXITER'] steps = eval(cfg['TRAIN']['STEPS']) batch_size = cfg['TRAIN']['BATCHSIZE'] subdivision = cfg['TRAIN']['SUBDIVISION'] ignore_thre = cfg['TRAIN']['IGNORETHRE'] random_resize = cfg['AUGMENTATION']['RANDRESIZE'] base_lr = cfg['TRAIN']['LR'] / batch_size / subdivision print('effective_batch_size = batch_size * iter_size = %d * %d' % (batch_size, subdivision)) # Learning rate setup def burnin_schedule(i): if i < burn_in: factor = pow(i / burn_in, 4) elif i < steps[0]: factor = 1.0 elif i < steps[1]: factor = 0.1 else: factor = 0.01 return factor # Initiate model model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre) if args.weights_path: print("loading darknet weights....", args.weights_path) parse_yolo_weights(model, args.weights_path) elif args.checkpoint: print("loading pytorch ckpt...", args.checkpoint) state = torch.load(args.checkpoint) if 'model_state_dict' in state.keys(): model.load_state_dict(state['model_state_dict']) else: model.load_state_dict(state) if cuda: print("using cuda") model = model.cuda() if args.tfboard: print("using tfboard") from tensorboardX import SummaryWriter tblogger = SummaryWriter(args.tfboard) model.train() imgsize = cfg['TRAIN']['IMGSIZE'] dataset = COCODataset(model_type=cfg['MODEL']['TYPE'], data_dir='COCO/', img_size=imgsize, augmentation=cfg['AUGMENTATION'], debug=args.debug) dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu) dataiterator = iter(dataloader) evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'], data_dir='COCO/', img_size=cfg['TEST']['IMGSIZE'], confthre=cfg['TEST']['CONFTHRE'], nmsthre=cfg['TEST']['NMSTHRE']) dtype = torch.cuda.FloatTensor if cuda else torch.FloatTensor # optimizer setup # set weight decay only on conv.weight params_dict = dict(model.named_parameters()) params = [] for key, value in params_dict.items(): if 'conv.weight' in key: params += [{'params':value, 'weight_decay':decay * batch_size * subdivision}] else: params += [{'params':value, 'weight_decay':0.0}] optimizer = optim.SGD(params, lr=base_lr, momentum=momentum, dampening=0, weight_decay=decay * batch_size * subdivision) iter_state = 0 if args.checkpoint: if 'optimizer_state_dict' in state.keys(): optimizer.load_state_dict(state['optimizer_state_dict']) iter_state = state['iter'] + 1 scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule) # start training loop for iter_i in range(iter_state, iter_size + 1): # COCO evaluation if iter_i % args.eval_interval == 0 and iter_i > 0: ap50_95, ap50 = evaluator.evaluate(model) model.train() if args.tfboard: tblogger.add_scalar('val/COCOAP50', ap50, iter_i) tblogger.add_scalar('val/COCOAP50_95', ap50_95, iter_i) # subdivision loop optimizer.zero_grad() for inner_iter_i in range(subdivision): try: imgs, targets, _, _ = next(dataiterator) # load a batch except StopIteration: dataiterator = iter(dataloader) imgs, targets, _, _ = next(dataiterator) # load a batch imgs = Variable(imgs.type(dtype)) targets = Variable(targets.type(dtype), requires_grad=False) loss = model(imgs, targets) loss.backward() optimizer.step() scheduler.step() if iter_i % 10 == 0: # logging current_lr = scheduler.get_lr()[0] * batch_size * subdivision print('[Iter %d/%d] [lr %f] ' '[Losses: xy %f, wh %f, conf %f, cls %f, total %f, imgsize %d]' % (iter_i, iter_size, current_lr, model.loss_dict['xy'], model.loss_dict['wh'], model.loss_dict['conf'], model.loss_dict['cls'], model.loss_dict['l2'], imgsize), flush=True) if args.tfboard: tblogger.add_scalar('train/total_loss', model.loss_dict['l2'], iter_i) # random resizing if random_resize: imgsize = (random.randint(0, 9) % 10 + 10) * 32 dataset.img_shape = (imgsize, imgsize) dataset.img_size = imgsize dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu) dataiterator = iter(dataloader) # save checkpoint if iter_i > 0 and (iter_i % args.checkpoint_interval == 0): torch.save({'iter': iter_i, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.checkpoint_dir, "snapshot"+str(iter_i)+".ckpt")) if args.tfboard: tblogger.close()
def main(): """ Visualize the detection result for the given image and the pre-trained model. """ parser = argparse.ArgumentParser() parser.add_argument('--gpu', type=int, default=0) parser.add_argument('--cfg', type=str, default='config/yolov3_default.cfg') parser.add_argument('--ckpt', type=str, help='path to the checkpoint file') parser.add_argument('--weights_path', type=str, default=None, help='path to weights file') parser.add_argument('--image', type=str) parser.add_argument( '--background', action='store_true', default=False, help='background(no-display mode. save "./output.png")') parser.add_argument('--detect_thresh', type=float, default=None, help='confidence threshold') args = parser.parse_args() with open(args.cfg, 'r') as f: cfg = yaml.load(f) imgsize = cfg['TEST']['IMGSIZE'] model = YOLOv3(cfg['MODEL']) confthre = cfg['TEST']['CONFTHRE'] nmsthre = cfg['TEST']['NMSTHRE'] if args.detect_thresh: confthre = args.detect_thresh img = cv2.imread(args.image) img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1)) img, info_img = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy) img = np.transpose(img / 255., (2, 0, 1)) img = torch.from_numpy(img).float().unsqueeze(0) if args.gpu >= 0: model.cuda(args.gpu) img = Variable(img.type(torch.cuda.FloatTensor)) else: img = Variable(img.type(torch.FloatTensor)) if args.weights_path: print("loading yolo weights %s" % (args.weights_path)) parse_yolo_weights(model, args.weights_path) else: print("loading checkpoint %s" % (args.ckpt)) model.load_state_dict(torch.load(args.ckpt)) model.eval() with torch.no_grad(): outputs = model(img) outputs = postprocess(outputs, 80, confthre, nmsthre) if outputs[0] is None: print("No Objects Deteted!!") return coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names( ) bboxes = list() classes = list() colors = list() for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]: cls_id = coco_class_ids[int(cls_pred)] print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred)) print('\t+ Label: %s, Conf: %.5f' % (coco_class_names[cls_id], cls_conf.item())) box = yolobox2label([y1, x1, y2, x2], info_img) bboxes.append(box) classes.append(cls_id) colors.append(coco_class_colors[int(cls_pred)]) if args.background: import matplotlib matplotlib.use('Agg') from utils.vis_bbox import vis_bbox import matplotlib.pyplot as plt vis_bbox(img_raw, bboxes, label=classes, label_names=coco_class_names, instance_colors=colors, linewidth=2) #plt.show() if args.background: plt.savefig('output.png')
def main(): args = parse_args() print("------------------------------------") print(" use {} dataset for demo. ".format(args.data)) print("------------------------------------") assert args.data in ['coco', 'drone'] if torch.cuda.is_available() and args.gpu >= 0: device = torch.device('cuda:{}'.format(args.gpu)) else: device = torch.device('cpu') # [TBM] gen n_classes from coco-format json file.. if args.data == 'coco': cfg_path = 'config/yolov3_default.cfg' n_classes = 80 if args.data == 'drone': cfg_path = 'config/yolov3_visdrone_default.cfg' n_classes = 10 with open(cfg_path, 'r') as f: cfg = yaml.load(f) imgsize = cfg['TEST']['IMGSIZE'] model = YOLOv3(n_classes=n_classes) confthre = cfg['TEST']['CONFTHRE'] nmsthre = cfg['TEST']['NMSTHRE'] if args.detect_thresh: confthre = args.detect_thresh img = cv2.imread(args.image) assert img is not None img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1)) img, info_img = preprocess(img, imgsize) # info = (h, w, nh, nw, dx, dy) img = torch.from_numpy(img).float().unsqueeze(0) model = model.to(device) img = Variable(img.to(device, dtype=torch.float32)) if args.weights_path: print("loading yolo weights %s" % (args.weights_path)) parse_yolo_weights(model, args.weights_path) else: print("loading checkpoint %s" % (args.ckpt)) model.load_state_dict(torch.load(args.ckpt)) model.eval() with torch.no_grad(): outputs = model(img) outputs = postprocess(outputs, n_classes, confthre, nmsthre) # [TBM] gen label_names from coco-format json file.. if args.data == 'coco': class_names, class_ids, class_colors = get_coco_label_names() if args.data == 'drone': class_names, class_ids, class_colors = get_visdrone_label_names() bboxes = list() classes = list() colors = list() for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]: cls_id = class_ids[int(cls_pred)] print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred)) print('\t+ Label: %s, Conf: %.5f' % (class_names[cls_id], cls_conf.item())) box = yolobox2label([y1, x1, y2, x2], info_img) bboxes.append(box) classes.append(cls_id) colors.append(class_colors[int(cls_pred)]) vis_bbox(img_raw, bboxes, label=classes, label_names=class_names, instance_colors=colors, linewidth=2) if args.window: plt.show() else: out_path = './output.png' plt.savefig(out_path, bbox_inches=0, pad_inches=0, dpi=100)
def main( image = None , gpu = -1, weights_path= f"{ Path(__file__).parent }/weights/yolov3.weights", background = False ): """ Visualize the detection result for the given image and the pre-trained model. """ print( weights_path ) my_path = Path( __file__ ).parent parser = argparse.ArgumentParser() parser.add_argument('--gpu', type=int, default= gpu ) parser.add_argument('--cfg', type=str, default=my_path/'config/yolov3_default.cfg') parser.add_argument('--ckpt', type=str, help='path to the checkpoint file') parser.add_argument('--weights_path', type=str, default= weights_path, help='path to weights file') parser.add_argument('--image', type=str , default= image ) parser.add_argument('--background', type=bool, default= background , help='background(no-display mode. save "./output.png")') parser.add_argument('--detect_thresh', type=float, default= 0.5 , help='confidence threshold') args = parser.parse_args() with open(args.cfg, 'r') as f: cfg = yaml.load(f) imgsize = cfg['TEST']['IMGSIZE'] model = YOLOv3(cfg['MODEL']) confthre = cfg['TEST']['CONFTHRE'] nmsthre = cfg['TEST']['NMSTHRE'] if args.detect_thresh: confthre = args.detect_thresh img = imread( args.image ) if img is None : print( "load image failed" ) print( args.image ) return img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1)) img, info_img = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy) img = np.transpose(img / 255., (2, 0, 1)) img = torch.from_numpy(img).float().unsqueeze(0) if args.gpu >= 0: model.cuda(args.gpu) img = Variable(img.type(torch.cuda.FloatTensor)) else: img = Variable(img.type(torch.FloatTensor)) assert args.weights_path or args.ckpt, 'One of --weights_path and --ckpt must be specified' if args.weights_path: print("loading yolo weights %s" % (args.weights_path)) parse_yolo_weights(model, args.weights_path) elif args.ckpt: print("loading checkpoint %s" % (args.ckpt)) state = torch.load(args.ckpt) if 'model_state_dict' in state.keys(): model.load_state_dict(state['model_state_dict']) else: model.load_state_dict(state) model.eval() with torch.no_grad(): outputs1 = model(img) # np.save("output.npy" , outputs.numpy() ) # torch.save( outputs1 , "outputs1.pt" ) out1 = torch.load( "outputs1.pt" ) rere = torch.equal( outputs1 , out1 ) outputs = postprocess(outputs1, 80, confthre, nmsthre) a = "hoho" if outputs[0] is None: print("No Objects Deteted!!") return coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names() bboxes = list() classes = list() colors = list() for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]: cls_id = coco_class_ids[int(cls_pred)] print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred)) print('\t+ Label: %s, Conf: %.5f' % (coco_class_names[cls_id], cls_conf.item())) box = yolobox2label([y1, x1, y2, x2], info_img) bboxes.append(box) classes.append(cls_id) colors.append(coco_class_colors[int(cls_pred)]) # args.background = True if args.background: import matplotlib matplotlib.use('Agg') from utils.vis_bbox import vis_bbox vis_bbox( img_raw, bboxes, label=classes, label_names=coco_class_names, instance_colors=colors, linewidth=2) if args.background: output = Path( "./output" ) output.mkdir( parents=True , exist_ok=True ) now = datetime.now().strftime("%Y-%m-%d %H-%M-%S") output /= f"output-{now}.png" plt.savefig( output ) return str( output.absolute() ) # return plt_to_qpixmap(plt.gca()) else : plt.show()
from utils.utils import * from utils.parse_yolo_weights import parse_yolo_weights from person_det import detect from mask_function import mask_plot with open('config/yolov3_default.cfg', 'r') as f: cfg = yaml.load(f) imgsize_y = cfg['TEST']['IMGSIZE'] model_Y = YOLOv3(cfg['MODEL']) confthre = cfg['TEST']['CONFTHRE'] nmsthre = cfg['TEST']['NMSTHRE'] model_Y.cuda(0) parse_yolo_weights(model_Y, 'weights/yolov3.weights') model_Y.eval() #################################################### mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(1, 1, 3) std = np.array([0.229, 0.224, 0.225], np.float32).reshape(1, 1, 3) mpii_edges = [[0, 1], [0, 2], [1, 3], [2, 4], [4, 6], [3, 5], [5, 6], [5, 7], [7, 9], [6, 8], [8, 10], [6, 12], [5, 11], [11, 12], [12, 14], [14, 16], [11, 13], [13, 15]] def show_2d(pts, c, edges,orig_image,min_x,min_y): for i in range(len(pts)): points = pts[i] num_joints = points.shape[0] points = ((points.reshape(num_joints, -1))).astype(np.int32) #for j in range(num_joints): # cv2.circle(img, (points[j, 0], points[j, 1]), 3, c, -1) for e in edges:
def main(): """ Visualize the detection result for the given image and the pre-trained model. """ parser = argparse.ArgumentParser() parser.add_argument('--gpu', type=int, default=0) args = parser.parse_args() with open('config/yolov3_default.cfg', 'r') as f: cfg = yaml.load(f) imgsize = cfg['TEST']['IMGSIZE'] model = YOLOv3(cfg['MODEL']) confthre = cfg['TEST']['CONFTHRE'] nmsthre = cfg['TEST']['NMSTHRE'] img = cv2.imread('data/mpii_87.png') orig_img = img.copy() img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1)) img, info_img = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy) img = np.transpose(img / 255., (2, 0, 1)) img = torch.from_numpy(img).float().unsqueeze(0) if args.gpu >= 0: model.cuda(args.gpu) img = Variable(img.type(torch.cuda.FloatTensor)) else: img = Variable(img.type(torch.FloatTensor)) parse_yolo_weights(model, 'weights/yolov3.weights') model.eval() with torch.no_grad(): outputs = model(img) outputs = postprocess(outputs, 80, confthre, nmsthre) if outputs[0] is None: print("No Objects Deteted!!") return coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names() bboxes = list() classes = list() colors = list() for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]: cls_id = coco_class_ids[int(cls_pred)] if cls_id==1: print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred)) print('\t+ Label: %s, Conf: %.5f' % (coco_class_names[cls_id], cls_conf.item())) box = yolobox2label([y1, x1, y2, x2], info_img) bboxes.append(box) y_min = int(box[0]) x_min = int(box[1]) y_max = int(box[2]) x_max = int(box[3]) #cropped = orig_img[y_min:y_max,x_min:x_max] cv2.rectangle(orig_img,(x_min,y_min),(x_max,y_max),(0,0,255),2) cv2.imshow('img',orig_img) cv2.waitKey(0)
def main(): """ Visualize the detection result for the given image and the pre-trained model. """ parser = argparse.ArgumentParser() parser.add_argument('--gpu', type=int, default=0) parser.add_argument('--cfg', type=str, default='config/yolov3_default_digestpath.cfg') parser.add_argument('--ckpt', type=str, help='path to the checkpoint file') parser.add_argument('--weights_path', type=str, default=None, help='path to weights file') parser.add_argument('--image', type=str) parser.add_argument( '--background', action='store_true', default=False, help='background(no-display mode. save "./output.png")') parser.add_argument('--detect_thresh', type=float, default=None, help='confidence threshold') parser.add_argument('--dataset', help='dataset to work with: {}'.format( Dataset.print_choices()), type=int, default=Dataset.SIGNET_RING) args = parser.parse_args() with open(args.cfg, 'r') as f: cfg = yaml.load(f) imgsize = cfg['TEST']['IMGSIZE'] model = YOLOv3(cfg['MODEL']) confthre = cfg['TEST']['CONFTHRE'] nmsthre = cfg['TEST']['NMSTHRE'] if args.detect_thresh: confthre = args.detect_thresh img = cv2.imread(args.image) img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1)) img, info_img = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy) img = np.transpose(img / 255., (2, 0, 1)) img = torch.from_numpy(img).float().unsqueeze(0) if args.gpu >= 0: model.cuda(args.gpu) img = Variable(img.type(torch.cuda.FloatTensor)) else: img = Variable(img.type(torch.FloatTensor)) assert args.weights_path or args.ckpt, 'One of --weights_path and --ckpt must be specified' if args.weights_path: print("loading yolo weights %s" % (args.weights_path)) parse_yolo_weights(model, args.weights_path) elif args.ckpt: print("loading checkpoint %s" % (args.ckpt)) state = torch.load(args.ckpt) if 'model_state_dict' in state.keys(): model.load_state_dict(state['model_state_dict']) else: model.load_state_dict(state) model.eval() with torch.no_grad(): outputs = model(img) outputs = postprocess(outputs, Dataset.NUM_CLASSES[args.dataset], confthre, nmsthre) if outputs[0] is None: print("No Objects Deteted!!") return bboxes = list() colors = list() for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]: print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred)) print('\t+ Conf: %.5f' % cls_conf.item()) box = yolobox2label([y1, x1, y2, x2], info_img) bboxes.append(box) colors.append(BOX_COLOR) if args.background: matplotlib.use('Agg') vis_bbox(img_raw, bboxes, instance_colors=colors, linewidth=2) plt.show() if args.background: plt.savefig('output.png')
def main(): """ YOLOv3 trainer. See README for details. """ args = parse_args() print("Setting Arguments.. : ", args) cuda = torch.cuda.is_available() and args.use_cuda os.makedirs(args.checkpoint_dir, exist_ok=True) # Parse config settings with open(args.cfg, 'r') as f: cfg = yaml.load(f) print("successfully loaded config file: ", cfg) momentum = cfg['TRAIN']['MOMENTUM'] decay = cfg['TRAIN']['DECAY'] burn_in = cfg['TRAIN']['BURN_IN'] iter_size = cfg['TRAIN']['MAXITER'] steps = eval(cfg['TRAIN']['STEPS']) batch_size = cfg['TRAIN']['BATCHSIZE'] subdivision = cfg['TRAIN']['SUBDIVISION'] ignore_thre = cfg['TRAIN']['IGNORETHRE'] random_resize = cfg['AUGMENTATION']['RANDRESIZE'] base_lr = cfg['TRAIN']['LR'] / batch_size / subdivision datatype = cfg['TRAIN']['DATATYPE'] print('effective_batch_size = batch_size * iter_size = %d * %d' % (batch_size, subdivision)) # Learning rate setup def burnin_schedule(i): if i < burn_in: factor = pow(i / burn_in, 4) elif i < steps[0]: factor = 1.0 elif i < steps[1]: factor = 0.1 else: factor = 0.01 return factor # Initiate model model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre) if args.weights_path: print("loading darknet weights....", args.weights_path) parse_yolo_weights(model, args.weights_path) elif args.checkpoint: print("loading pytorch ckpt...", args.checkpoint) state = torch.load(args.checkpoint) if 'model_state_dict' in state.keys(): model.load_state_dict(state['model_state_dict']) else: model.load_state_dict(state) if cuda: print("using cuda") model = model.cuda() if args.tfboard: print("using tfboard") from tensorboardX import SummaryWriter tblogger = SummaryWriter(args.tfboard) model.train() coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names() imgsize = cfg['TRAIN']['IMGSIZE'] if datatype=='voc': dataset = VOCDataset(model_type=cfg['MODEL']['TYPE'], data_dir='../../VOCdevkit/VOC2007', img_size=imgsize, augmentation=cfg['AUGMENTATION'], debug=args.debug) print('load voc dataset successfully') else: dataset = COCODataset(model_type=cfg['MODEL']['TYPE'], data_dir='COCO/', img_size=imgsize, augmentation=cfg['AUGMENTATION'], debug=args.debug) print('load COCO dataset successfully') evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'], data_dir='COCO/', img_size=cfg['TEST']['IMGSIZE'], confthre=cfg['TEST']['CONFTHRE'], nmsthre=cfg['TEST']['NMSTHRE']) dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu) dataiterator = iter(dataloader) dtype = torch.cuda.FloatTensor if cuda else torch.FloatTensor # optimizer setup # set weight decay only on conv.weight params_dict = dict(model.named_parameters()) params = [] for key, value in params_dict.items(): if 'conv.weight' in key: params += [{'params': value, 'weight_decay': decay * batch_size * subdivision}] else: params += [{'params': value, 'weight_decay': 0.0}] optimizer = optim.SGD(params, lr=base_lr, momentum=momentum, dampening=0, weight_decay=decay * batch_size * subdivision) iter_state = 0 if args.checkpoint: if 'optimizer_state_dict' in state.keys(): optimizer.load_state_dict(state['optimizer_state_dict']) iter_state = state['iter'] + 1 #学习率控制 Sets the learning rate of each parameter group to the initial lr times a given function. When last_epoch=-1, sets initial lr as lr. scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule) # result=evals(model) # print(result) # start training loop # print('args.eval_interval',args.eval_interval) for iter_i in range(iter_state, iter_size + 1): if iter_i % (args.eval_interval*2) == 0 and iter_i > 0: if datatype=='voc': result=evals(model) print(result) else: ap50_95, ap50 = evaluator.evaluate(model) print(ap50_95, ap50) model.train() if args.tfboard: tblogger.add_scalar('val/COCOAP50', ap50, iter_i) tblogger.add_scalar('val/COCOAP50_95', ap50_95, iter_i) if iter_i % (40000) == 0 and iter_i > 0: draw(model,datatype,imgsize) # subdivision loop optimizer.zero_grad() for inner_iter_i in range(subdivision): try: imgs, targets, info_img, id_ = next(dataiterator) # load a batch except StopIteration: dataiterator = iter(dataloader) imgs, targets, info_img, id_ = next(dataiterator) # load a batch imgs = Variable(imgs.type(dtype)) targets = Variable(targets.type(dtype), requires_grad=False) loss = model(imgs, targets) loss.backward() optimizer.step() scheduler.step() if iter_i % 10 == 0: # logging current_lr = scheduler.get_lr()[0] * batch_size * subdivision print('[Iter %d/%d] [lr %f] ' '[Losses: xy %f, wh %f, conf %f, cls %f, total %f, imgsize %d]' % (iter_i, iter_size, current_lr, model.loss_dict['xy'], model.loss_dict['wh'], model.loss_dict['conf'], model.loss_dict['cls'], model.loss_dict['l2'], imgsize)) if args.tfboard: tblogger.add_scalar('train/total_loss', model.loss_dict['l2'], iter_i) # random resizing # 变输入大小,利用了yolov3网络的全卷积,使得模型不受图像大小的改变而影响参数。 if random_resize: imgsize = (random.randint(0, 9) % 10 + 10) * 32 dataset.img_shape = (imgsize, imgsize) dataset.img_size = imgsize dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu) dataiterator = iter(dataloader) if iter_i % 100 == 0: model.eval() if datatype=='voc': img_file = os.path.join('../../VOCdevkit/VOC2007', 'JPEGImages', id_[0] + '.jpg') else: img_file = os.path.join('COCO', 'train2017', '{:012}'.format(id_) + '.jpg') img = cv2.imread(img_file) img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1)) # print(img_raw.shape) # print(img_raw) # print(imgs) img, info_img_t = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy) img = np.transpose(img / 255., (2, 0, 1)) img = torch.from_numpy(img).float().unsqueeze(0) img = Variable(img.type(torch.cuda.FloatTensor)) outputs = model(img) #outputs.shape : torch.Size([1, 12348, 85]) outputs = postprocess(outputs, 80, 0.5, 0.5) # imgs.shape : torch.Size([1, 3, 608, 608]) # outputs[0].shape :torch.Size([3, 7]) # targets.shape :torch.Size([1, 50, 5]) # print(outputs) if outputs[0] is not None: bboxes = list() classes = list() colors = list() # print(info_img_t) info_img=tuple(info_img) # print(info_img) for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]: cls_id = coco_class_ids[int(cls_pred)] # print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred)) # print('\t+ Label: %s, Conf: %.5f' % # (coco_class_names[cls_id], cls_conf.item())) # print([y1, x1, y2, x2]) box = yolobox2label([y1, x1, y2, x2], info_img_t) bboxes.append(box) classes.append(cls_id) colors.append(coco_class_colors[int(cls_pred)]) vis_bbox( img_raw, bboxes, label=classes, label_names=coco_class_names, instance_colors=colors, linewidth=2) plt.savefig('output/'+str(iter_i)+'.jpg') model.train() # save checkpoint if iter_i > 0 and (iter_i % args.checkpoint_interval == 0): torch.save({'iter': iter_i, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(args.checkpoint_dir, "snapshot" + str(iter_i) + ".ckpt")) if args.tfboard: tblogger.close()
def main(): args = parse_args() print("------------------------------------") print(" use {} dataset for demo. ".format(args.data)) print("------------------------------------") assert args.data in ['coco', 'drone'] if torch.cuda.is_available() and args.gpu >= 0: device = torch.device('cuda:{}'.format(args.gpu)) else: device = torch.device('cpu') # [TBM] gen n_classes from coco-format json file.. if args.data == 'coco': cfg_path = 'config/yolov3_default.cfg' n_classes = 80 if args.data == 'drone': cfg_path = 'config/yolov3_visdrone_default.cfg' n_classes = 10 with open(cfg_path, 'r') as f: cfg = yaml.load(f) imgsize = cfg['TEST']['IMGSIZE'] model = YOLOv3(n_classes=n_classes) confthre = cfg['TEST']['CONFTHRE'] nmsthre = cfg['TEST']['NMSTHRE'] if args.detect_thresh: confthre = args.detect_thresh model = model.to(device) if args.weights_path: print("loading yolo weights %s" % (args.weights_path)) parse_yolo_weights(model, args.weights_path) else: print("loading checkpoint %s" % (args.ckpt)) model.load_state_dict(torch.load(args.ckpt)) model.eval() dir_name = os.path.basename(os.path.dirname(args.in_dir + '/')) out_dir = os.path.join(args.out_dir, dir_name) os.makedirs(out_dir, exist_ok=True) img_files = os.listdir(args.in_dir) img_files.sort() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) for i in range(0, len(img_files), args.step): filename = img_files[i] img_path = os.path.join(args.in_dir, filename) img = cv2.imread(img_path) assert img is not None start.record() img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1)) img, info_img = preprocess(img, imgsize) # info = (h, w, nh, nw, dx, dy) img = torch.from_numpy(img).float().unsqueeze(0) img = Variable(img.to(device, dtype=torch.float32)) with torch.no_grad(): outputs = model(img) outputs = postprocess(outputs, n_classes, confthre, nmsthre) end.record() torch.cuda.synchronize() # [TBM] gen label_names from coco-format json file.. if args.data == 'coco': class_names, class_ids, class_colors = get_coco_label_names() if args.data == 'drone': class_names, class_ids, class_colors = get_visdrone_label_names() bboxes, classes, colors = list(), list(), list() if outputs[0] is None: outputs[0] = [] if args.verbose: print("=====================================") print("{}, {:.2f} [fps]".format(filename, 1000.0 / start.elapsed_time(end))) for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]: cls_id = class_ids[int(cls_pred)] if args.verbose: print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred)) print('\t+ Label: %s, Conf: %.5f' % (class_names[cls_id], cls_conf.item())) box = yolobox2label([y1, x1, y2, x2], info_img) bboxes.append(box) classes.append(cls_id) colors.append(class_colors[int(cls_pred)]) if args.verbose: print() vis_bbox(img_raw, bboxes, label=classes, label_names=class_names, instance_colors=colors, linewidth=2) basename, _ = os.path.splitext(filename) out_path = os.path.join(out_dir, '{}.png'.format(basename)) plt.savefig(out_path, bbox_inches=0, pad_inches=0, dpi=100) plt.close() print("Done!")