def __init__(self, bgr=True, gpu_device=0, **kwargs): self.__dict__.update(self._defaults) # set up default values # for portability between keras-yolo3/yolo.py and this if 'model_path' in kwargs: kwargs['weights'] = kwargs['model_path'] if 'score' in kwargs: kwargs['thresh'] = kwargs['score'] self.__dict__.update(kwargs) # update with user overrides self.class_names = self._get_class() # self.model = Darknet(self.config) # self.model.load_weights(self.weights) self.model = Yolov4(n_classes=len(self.class_names), inference=True) checkpoint = torch.load(self.weights, map_location=torch.device('cpu')) # checkpoint = self._rename_checkpoint(checkpoint) self.model.load_state_dict(checkpoint) self.device = gpu_device self.model.cuda(self.device) self.model.eval() self.bgr = bgr if self.half: self.model.half() # warm up self._detect([np.zeros((10, 10, 3), dtype=np.uint8)]) # self._detect([np.zeros((10,10,3), dtype=np.uint8), np.zeros((10,10,3), dtype=np.uint8), np.zeros((10,10,3), dtype=np.uint8), np.zeros((10,10,3), dtype=np.uint8)]) print('Warmed up!')
def transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W): model = Yolov4(n_classes=n_classes, inference=True) pretrained_dict = torch.load(weight_file, map_location=torch.device('cuda')) model.load_state_dict(pretrained_dict) x = torch.randn((batch_size, 3, IN_IMAGE_H, IN_IMAGE_W), requires_grad=True) # .cuda() onnx_file_name = "yolov4_{}_3_{}_{}.onnx".format(batch_size, IN_IMAGE_H, IN_IMAGE_W) # Export the model print('Export the onnx model ...') torch.onnx.export(model, x, onnx_file_name, export_params=True, opset_version=11, do_constant_folding=True, input_names=['input'], output_names=['boxes', 'confs'], dynamic_axes=None) print('Onnx model exporting done') return onnx_file_name
def detect_cv2(namesfile,anchors,weightfile,imgfile): import cv2 class_names = load_class_names(namesfile) n_classes=len(class_names) m = Yolov4(anchors=anchors,yolov4conv137weight=None, n_classes=n_classes, inference=True) pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) m.load_state_dict(pretrained_dict) print('Loading weights from %s... Done!' % (weightfile)) if use_cuda: m.cuda() img = cv2.imread(imgfile) sized = cv2.resize(img, (416, 416)) sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB) for i in range(2): start = time.time() boxes = do_detect(m, sized, 0.4, 0.6, use_cuda) finish = time.time() if i == 1: print('%s: Predicted in %f seconds.' % (imgfile, (finish - start))) plot_boxes_cv2(img, boxes[0], savename='predictions.jpg', class_names=class_names)
def load_model_weight(model_path, num_classes): model = Yolov4(n_classes=num_classes) print("------------loading model") model.load_state_dict( torch.load(model_path, map_location=torch.device('cuda'))) print("------------done loading") return model
def __init__(self, weightfile, use_cuda=True): self.use_cuda = use_cuda self.model = Yolov4(yolov4conv137weight=None, n_classes=80, inference=True) pretrained_dict = torch.load( weightfile, map_location=torch.device("cuda" if use_cuda else "cpu") ) self.model.load_state_dict(pretrained_dict) if self.use_cuda: self.model.cuda()
def __init__(self, weightfile, use_cuda=True): if use_cuda and not torch.cuda.is_available(): raise Exception( "Selected use_cuda=True, but cuda is not available to Pytorch" ) self.use_cuda = use_cuda self.model = Yolov4(yolov4conv137weight=None, n_classes=80, inference=True) pretrained_dict = torch.load( weightfile, map_location=torch.device("cuda" if use_cuda else "cpu") ) self.model.load_state_dict(pretrained_dict) if self.use_cuda: self.model.cuda()
def transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W): model = Yolov4(n_classes=n_classes, inference=True) pretrained_dict = torch.load(weight_file, map_location=torch.device('cuda')) model.load_state_dict(pretrained_dict) input_names = ["input"] output_names = ['boxes', 'confs'] dynamic = False if batch_size <= 0: dynamic = True if dynamic: x = torch.randn((1, 3, IN_IMAGE_H, IN_IMAGE_W), requires_grad=True) onnx_file_name = "yolov4_-1_3_{}_{}_dynamic.onnx".format(IN_IMAGE_H, IN_IMAGE_W) dynamic_axes = {"input": {0: "batch_size"}, "boxes": {0: "batch_size"}, "confs": {0: "batch_size"}} # Export the model print('Export the onnx model ...') torch.onnx.export(model, x, onnx_file_name, export_params=True, opset_version=11, do_constant_folding=True, input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes) print('Onnx model exporting done') return onnx_file_name else: x = torch.randn((batch_size, 3, IN_IMAGE_H, IN_IMAGE_W), requires_grad=True) onnx_file_name = "yolov4_{}_3_{}_{}_static.onnx".format(batch_size, IN_IMAGE_H, IN_IMAGE_W) # Export the model print('Export the onnx model ...') torch.onnx.export(model, x, onnx_file_name, export_params=True, opset_version=11, do_constant_folding=True, input_names=input_names, output_names=output_names, dynamic_axes=None) print('Onnx model exporting done') return onnx_file_name
def transform_to_onnx(weight_file, onnx_file_name, batch_size, n_classes, input_h, input_w): model = Yolov4(n_classes=n_classes, inference=True).to(device) pretrained_dict = torch.load(weight_file) model.load_state_dict(pretrained_dict) input_names = ["input"] output_names = ["boxes", "confs"] x = torch.randn(batch_size, 3, input_h, input_w).to(device) print('Export the onnx model ...') torch.onnx.export(model, x, onnx_file_name, export_params=True, opset_version=11, do_constant_folding=True, input_names=input_names, output_names=output_names, dynamic_axes=None) print('Onnx model exporting done')
def __init__(self,weightfile,num_classes=80,width=416,height=416,use_cuda = True,is_half=True): self.num_classes = num_classes self.model = Yolov4(yolov4conv137weight=None, n_classes=self.num_classes, inference=True) pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) self.model.load_state_dict(pretrained_dict) self.use_cuda = use_cuda self.is_half=is_half self.width=width self.height=height if self.use_cuda: if self.is_half: self.model=self.model.cuda().half() else: self.model=self.model.cuda() if self.num_classes == 20: namesfile = 'data/voc.names' elif self.num_classes == 80: namesfile = 'data/coco.names' else: namesfile = 'data/x.names' self.class_names = load_class_names(namesfile) print("initialize_a_pytorch")
def detect_cv2_camera(namesfile,anchors,weightfile,size=416): import cv2 anchors=np.array(anchors)*size anchors=anchors.tolist() class_names = load_class_names(namesfile) n_classes=len(class_names) m = Yolov4(anchors=anchors,yolov4conv137weight=None, n_classes=n_classes, inference=True) pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) m.load_state_dict(pretrained_dict) print('Loading weights from %s... Done!' % (weightfile)) if use_cuda: m.cuda() # cap = cv2.VideoCapture(0) cap = cv2.VideoCapture("../dataset/sample1.mov") cap.set(3, 1280) cap.set(4, 720) print("Starting the YOLO loop...") while True: ret, img = cap.read() if(ret==False): continue sized = cv2.resize(img, (size,size)) sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB) start = time.time() boxes = do_detect(m, sized, 0.2, 0.6, use_cuda) finish = time.time() print('Predicted in %f seconds.' % (finish - start)) result_img = plot_boxes_cv2(img, boxes[0], savename=None, class_names=class_names) cv2.imshow('Yolo demo', result_img) cv2.waitKey(1) cap.release()
def init_net(pretrained, n_classes): # m = Darknet(cfgfile) # model = Yolov4(cfg.pretrained, n_classes=cfg.classes) m = Yolov4(n_classes=n_classes) # model = Yolov4(n_classes=n_classes) pretrained_dict = torch.load(pretrained, map_location=torch.device('cuda')) from collections import OrderedDict new_state_dict = OrderedDict() for key, value in pretrained_dict.items(): new_key = key.replace('module.', '') new_state_dict[new_key] = value m.load_state_dict(new_state_dict) # m.print_network() # m.load_weights(weightfile) print('Loading weights from %s... Done!' % (pretrained)) if use_cuda: m.cuda() return m
def get_image_objects(image: Image): models_path = os.path.join(os.getcwd(), 'modules', 'pytorch-YOLOv4') if models_path not in sys.path: sys.path.insert(1, models_path) from models import Yolov4 from tool.utils import load_class_names device = 'cpu' model = Yolov4(yolov4conv137weight=None, n_classes=80, inference=True) model.load_state_dict( torch.load(os.path.join(os.getcwd(), 'models', 'yolov4', 'yolov4.pth'), map_location=torch.device(device))) model.eval() model = model.to(torch.device(device)) image_tensor = ToTensor()(image) image_tensor = image_tensor.unsqueeze(0) image_tensor = image_tensor.to(torch.device(device)) with torch.no_grad(): output = model(image_tensor) conf_thresh = 0.4 nms_thresh = 0.6 boxes = post_process_boxes(conf_thresh, nms_thresh, output) class_names = load_class_names( os.path.join(models_path, 'data', 'coco.names')) return list( map( lambda obj: { "x1": obj[0], "y1": obj[1], "x2": obj[2], "y2": obj[3], "class_name": class_names[obj[6]] }, boxes[0]))
def detect_skimage(namesfile,anchors,weightfile, imgfile): from skimage import io from skimage.transform import resize class_names = load_class_names(namesfile) n_classes=len(class_names) m = Yolov4(anchors=anchors,yolov4conv137weight=None, n_classes=n_classes, inference=True) pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) m.load_state_dict(pretrained_dict) print('Loading weights from %s... Done!' % (weightfile)) if use_cuda: m.cuda() img = io.imread(imgfile) sized = resize(img, (m.width, m.height)) * 255 for i in range(2): start = time.time() boxes = do_detect(m, sized, 0.4, 0.4, use_cuda) finish = time.time() if i == 1: print('%s: Predicted in %f seconds.' % (imgfile, (finish - start))) plot_boxes_cv2(img, boxes, savename='predictions.jpg', class_names=class_names)
def train_yolov4(cfg): logging = init_logger(log_dir=cfg.TRAIN_TENSORBOARD_DIR) os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') if cfg.use_darknet_cfg: model = Darknet(cfg.cfgfile) else: model = Yolov4(cfg.pretrained, n_classes=cfg.classes) #TODO, load checkpoints ''' WORK_FOLDER = '/mnt/bos/modules/perception/emergency_detection' weightfile = os.path.join(WORK_FOLDER, 'checkpoints/Yolov4_epoch291.pth') pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) model.load_state_dict(pretrained_dict) ''' if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) model.to(device=device) try: train(model=model, config=cfg, epochs=cfg.TRAIN_EPOCHS, device=device, ) except KeyboardInterrupt: torch.save(model.state_dict(), 'INTERRUPTED.pth') logging.info('Saved interrupt') try: sys.exit(0) except SystemExit: os._exit(0)
return logging def _get_date_str(): now = datetime.datetime.now() return now.strftime('%Y-%m-%d_%H-%M') if __name__ == "__main__": logging = init_logger(log_dir='log') cfg = Cfg device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') model = Yolov4("yolov4.conv.137.pth", n_classes=cfg.classes) model.to(device=device) try: train(model=model, config=cfg, epochs=cfg.TRAIN_EPOCHS, device=device, ) except KeyboardInterrupt: torch.save(model.state_dict(), 'INTERRUPTED.pth') logging.info('Saved interrupt') try: sys.exit(0) except SystemExit: os._exit(0)
def plot_lines(image, boxes): width = img.shape[1] height = img.shape[0] angled = np.sqrt(width**2 + height**2) for box in boxes: x = (box[2] + box[0]) / 2 * width y = (box[3] + box[1]) / 2 * height print() return image if __name__ == '__main__': model = Yolov4(None, n_classes=80, inference=True) state = torch.load('16.pth') model.load_state_dict(state['model']) del state model.eval() rawimg = cv2.imread( '/media/palm/data/coco/images/val2017/000000289343.jpg') rawimg = cv2.cvtColor(rawimg, cv2.COLOR_BGR2RGB) rawimg, bboxes = resize_image(rawimg, None, 608, 608) img = rawimg.copy().transpose(2, 0, 1) inputs = torch.from_numpy(np.expand_dims(img, 0).astype('float32')).div(255.0) output = model(inputs) boxes = post_processing(img, 0.4, 0.4, output)
is_first_frame = True num_frame = 0 counter_detection = 0 person_id = 0 prev_id = 0 prev_coords = None people_actions = dict() person_inside = dict() person_outside = dict() people_actions['in'] = 0 people_actions['out'] = 0 check_gpu() model = Yolov4(weight_path='yolo4_config/yolov4.weights', class_name_path='yolo4_config/coco_classes.txt') videofile = cv2.VideoCapture(video_path) fps = videofile.get(cv2.CAP_PROP_FPS) fourcc = cv2.VideoWriter_fourcc(*'MP4V') output_video = cv2.VideoWriter(output_name, fourcc, fps, (640, 480)) rr, first_frame = videofile.read() if initialize_door_by_yourself: door_array = select_object(first_frame)[0] else: door_array = [361, 20, 507, 352] while rr: ret, frame = videofile.read() if not ret:
logging.getLogger('').addHandler(console) return logging if __name__ == "__main__": logging = init_logger(log_dir='log') cfg = get_args(**Cfg) os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') if cfg.use_darknet_cfg: model = Darknet(cfg.cfgfile) else: model = Yolov4(cfg.weight, n_classes=cfg.classes) # pretrained_dict = torch.load(cfg.load) # model_dict = model.state_dict() # pretrained_dict = {k: v for k, v in pretrained_dict.items() if # k in model_dict} # pretrained_dict只保留了model_dict中存在的键。为什么直接load它会报错。要先给model_dict更新。 # model_dict.update(pretrained_dict) # model.load_state_dict(model_dict) total_params = sum(p.numel() for p in model.parameters()) print(f'{total_params:,} total parameters.') total_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f'{total_trainable_params:,} training parameters.') if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model)
def main(): n_classes = 18 weightfile = r'D:\DeepBlueProjects\chem-lab\pytorch-YOLOv4-master\checkpoints\Yolov4_epoch6.pth' imgfile = r'D:\Data\CMS01_single-end\val\JPEGImages\frontfront_0518.jpg' base_dir = r'D:\Data\chem-yolov4\eval-dataset\top' gt_path = os.path.join(base_dir, 'gt.json') name_id_path = os.path.join(base_dir, 'name_id.json') with open(gt_path, 'r') as f: gt_dict = json.load(f) with open(name_id_path, 'r') as f: name_id_dict = json.load(f) input_size = (960, 960) model = Yolov4(yolov4conv137weight=None, n_classes=n_classes, inference=True) pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) from collections import OrderedDict new_state_dict = OrderedDict() for k, v in pretrained_dict.items(): name = k[7:] # remove `module.`,表面从第7个key值字符取到最后一个字符,正好去掉了module. new_state_dict[name] = v # 新字典的key值对应的value为一一对应的值。 model.load_state_dict(pretrained_dict) use_cuda = True if use_cuda: model.cuda() data_txt = os.path.join(base_dir, 'data.txt') save_dir = os.path.join(base_dir, 'JPEGImages_pred') result_dir = os.path.join(base_dir, 'result_txt') with open(data_txt, 'r') as f: imgfiles = f.readlines() box_list = [] for imgfile in imgfiles: img = cv2.imread(imgfile.strip('\n')) img_h, img_w, _ = img.shape img_name = imgfile.split('\\')[-1].strip('\n') img_id = name_id_dict[img_name] result_txt = os.path.join(result_dir, img_name[:-4] + '.txt') result_f = open(result_txt, 'w') # Inference input size is 416*416 does not mean training size is the same # Training size could be 608*608 or even other sizes # Optional inference sizes: # Hight in {320, 416, 512, 608, ... 320 + 96 * n} # Width in {320, 416, 512, 608, ... 320 + 96 * m} sized = cv2.resize(img, (input_size[1], input_size[0])) sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB) # for i in range(2): # This 'for' loop is for speed check # Because the first iteration is usually longer boxes = do_detect(model, sized, 0.01, 0.3, use_cuda) for box in boxes[0]: x1 = int((box[0] - box[2] / 2.0) * img_w) y1 = int((box[1] - box[3] / 2.0) * img_h) x2 = int((box[0] + box[2] / 2.0) * img_w) y2 = int((box[1] + box[3] / 2.0) * img_h) w = x2 - x1 h = y2 - y1 if len(box) >= 7: cls_conf = box[5] cls_id = box[6] box_list.append({ "image_id": img_id, "category_id": int(cls_id), "bbox": [x1, y1, w, h], "score": float(cls_conf) }) string = ','.join([ str(cls_id), str(x1), str(y1), str(x2), str(y2), str(cls_conf) ]) + '\n' result_f.write(string) if cls_conf > 0.3: cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 255), 1) cv2.putText(img, str(cls_id), (int(x1 + 10), int(y1 + 20)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 1) cv2.putText(img, str(round(cls_conf, 3)), (int(x1 + 30), int(y1 + 20)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (20, 240, 0), 1) else: print('????') result_f.close() namesfile = 'data/chem.names' class_names = load_class_names(namesfile) save_name = os.path.join(save_dir, img_name) plot_boxes_cv2(img, boxes[0], save_name, class_names) # cv2.imshow('result', img) # cv2.waitKey(0) # cv2.destroyAllWindows() info, map_iou0_5 = get_coco_mAP(gt_dict, box_list) # print("---base_eval---epoch%d"%real_epoch) print(info)
n_classes = int(sys.argv[1]) namesfile = sys.argv[2] dataset_dir = sys.argv[3] width = int(sys.argv[4]) height = int(sys.argv[5]) weightfile = sys.argv[6] eval_model = True else: print('Usage: ') print( ' python models.py [num_classes] [names_file] [datasetdir] [width] [height] ([pretrained])' ) if eval_model: model = Yolov4(yolov4conv137weight=None, n_classes=n_classes, inference=True) pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) model.load_state_dict(pretrained_dict) use_device = 1 # -1 for cpu, 0,1,2,... for gpu 0,1,2,... if use_device >= 0: # model.cuda() model.to("cuda:{}".format(use_device)) if not output_video: window_name = "gt vs. pred" if eval_model else "gt" cv2.namedWindow(window_name, 0) cv2.resizeWindow(window_name, 1000, 500)
# -*- coding: utf-8 -*- """ Created on Wed Feb 3 16:24:19 2021 @author: vinic_000 """ from models import Yolov4 model = Yolov4(weight_path='yolov4.weights', class_name_path='class_names/coco_classes.txt') model.predict('img/pessoas.png') model.predict('img/vaca.png') model.predict('img/cavalo.png') model.predict('img/escola.jpg') model.predict('img/kite.jpg')
def _get_date_str(): now = datetime.datetime.now() return now.strftime('%Y-%m-%d_%H-%M') if __name__ == "__main__": logging = init_logger(log_dir='log') cfg = get_args(**Cfg) os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') if cfg.use_darknet_cfg: model = Darknet(cfg.cfgfile) else: model = Yolov4(cfg.pretrained, n_classes=cfg.classes) if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) model.to(device=device) try: train( model=model, config=cfg, epochs=cfg.TRAIN_EPOCHS, device=device, ) except KeyboardInterrupt: torch.save(model.state_dict(), 'INTERRUPTED.pth') logging.info('Saved interrupt')
def train(model, device, config, epochs=5, batch_size=1, save_cp=True, log_step=20, img_scale=0.5, freeze_backbone=False): """ Train the YOLOv4 network with given configurations """ train_dataset = Yolo_dataset(config.train_label, config, train=True) val_dataset = Yolo_dataset(config.val_label, config, train=False) n_train = len(train_dataset) n_val = len(val_dataset) train_loader = DataLoader(train_dataset, batch_size=config.batch // config.subdivisions, shuffle=True, num_workers=8, pin_memory=True, drop_last=True, collate_fn=collate) val_loader = DataLoader(val_dataset, batch_size=config.batch // config.subdivisions, shuffle=True, num_workers=8, pin_memory=True, drop_last=True, collate_fn=val_collate) global_step = 0 logging.info(f'''Starting training: Epochs: {epochs} Batch size: {config.batch} Subdivisions: {config.subdivisions} Learning rate: {config.learning_rate} Training size: {n_train} Validation size: {n_val} Checkpoints: {save_cp} Device: {device.type} Images size: {config.width} Optimizer: {config.TRAIN_OPTIMIZER} Dataset classes: {config.classes} Train label path:{config.train_label} Pretrained: {config.pretrained} ''') # Learning rate setup def burnin_schedule(i): if i < config.burn_in: factor = pow(i / config.burn_in, 4) elif i < config.steps[0]: factor = 1.0 elif i < config.steps[1]: factor = 0.1 else: factor = 0.01 return factor # Optimiser alternatives if config.TRAIN_OPTIMIZER.lower() == 'adam': optimizer = optim.Adam( model.parameters(), lr=config.learning_rate / config.batch, betas=(0.9, 0.999), eps=1e-08, ) elif config.TRAIN_OPTIMIZER.lower() == 'sgd': optimizer = optim.SGD( params=model.parameters(), lr=config.learning_rate / config.batch, momentum=config.momentum, weight_decay=config.decay, ) scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule) criterion = Yolo_loss(device=device, batch=config.batch // config.subdivisions, n_classes=config.classes) save_prefix = 'Yolov4_epoch' saved_models = deque() # Set model to training mode model.train() for epoch in range(epochs): epoch_loss = 0 epoch_step = 0 # Freeze backbone first epoch if freeze_backbone and epoch < 2: for name, p in model.named_parameters(): # Freeze everything but the head if not 'head' in name.split('.')[0]: p.requires_grad = False if (epoch == 0) else True with tqdm(total=n_train, desc=f'Epoch {epoch + 1}/{epochs}', unit='img', ncols=200) as progress_bar: for batch in train_loader: global_step += 1 epoch_step += 1 images = batch[0] bboxes = batch[1] images = images.to(device=device, dtype=torch.float32) bboxes = bboxes.to(device=device) bboxes_pred = model(images) loss, loss_xy, loss_wh, loss_obj, loss_cls, loss_l2 = criterion( bboxes_pred, bboxes) loss.backward() epoch_loss += loss.item() if global_step % config.subdivisions == 0: optimizer.step() scheduler.step() model.zero_grad() if global_step % (log_step * config.subdivisions) == 0: progress_bar.set_postfix( **{ 'loss (batch)': loss.item(), 'loss_xy': loss_xy.item(), 'loss_wh': loss_wh.item(), 'loss_obj': loss_obj.item(), 'loss_cls': loss_cls.item(), 'loss_l2': loss_l2.item(), 'lr': scheduler.get_lr()[0] * config.batch }) logging.debug( 'Train step_{}: loss : {}, loss xy : {}, loss wh : {},' ' loss obj : {}, loss cls : {}, loss l2 : {}, lr : {}'. format(global_step, loss.item(), loss_xy.item(), loss_wh.item(), loss_obj.item(), loss_cls.item(), loss_l2.item(), scheduler.get_lr()[0] * config.batch)) progress_bar.update(images.shape[0]) eval_model = Yolov4(cfg.pretrained, n_classes=cfg.classes, inference=True) if torch.cuda.device_count() > 1: eval_model.load_state_dict(model.module.state_dict()) else: eval_model.load_state_dict(model.state_dict()) eval_model.to(device) """ evaluator = evaluate(eval_model, val_loader, config, device) del eval_model stats = evaluator.coco_eval['bbox'].stats writer.add_scalar('train/AP', stats[0], global_step) writer.add_scalar('train/AP50', stats[1], global_step) writer.add_scalar('train/AP75', stats[2], global_step) writer.add_scalar('train/AP_small', stats[3], global_step) writer.add_scalar('train/AP_medium', stats[4], global_step) writer.add_scalar('train/AP_large', stats[5], global_step) writer.add_scalar('train/AR1', stats[6], global_step) writer.add_scalar('train/AR10', stats[7], global_step) writer.add_scalar('train/AR100', stats[8], global_step) writer.add_scalar('train/AR_small', stats[9], global_step) writer.add_scalar('train/AR_medium', stats[10], global_step) writer.add_scalar('train/AR_large', stats[11], global_step) """ # Save model to file if save_cp: try: os.makedirs(config.checkpoints, exist_ok=True) logging.info('Created checkpoint directory') except OSError: pass save_path = os.path.join(config.checkpoints, f'{save_prefix}{epoch + 1}.pth') torch.save(model.state_dict(), save_path) logging.info(f'Checkpoint {epoch + 1} saved !') saved_models.append(save_path) if len(saved_models) > config.keep_checkpoint_max > 0: model_to_remove = saved_models.popleft() try: os.remove(model_to_remove) except: logging.info(f'failed to remove {model_to_remove}')
import tvm from tvm import relay import numpy as np from tvm.contrib.download import download_testdata import torch from tvm.contrib import graph_runtime from PIL import Image from torchvision import transforms from models import Yolov4 model = Yolov4(inference=True) model.eval() input_shape = [1, 3, 512, 512] input_data = torch.randn(input_shape) scripted_model = torch.jit.trace(model, input_data).eval() img_url = "https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true" img_path = download_testdata(img_url, "cat.png", module="data") img = Image.open(img_path).resize((512, 512)) my_preprocess = transforms.Compose( [ transforms.Resize(512), transforms.CenterCrop(512), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ) img = my_preprocess(img) img = np.expand_dims(img, 0) input_name = "input0"
def evaluate_all(device=None): """ """ device = device or torch.device("cuda") print(f"device = {device}") val_dataset = ACNE04(label_path=Cfg.val_label, cfg=Cfg, train=False) val_loader = DataLoader( dataset=val_dataset, batch_size=Cfg.batch // Cfg.subdivisions, shuffle=True, num_workers=8, pin_memory=True, drop_last=True, collate_fn=val_collate, ) coco = convert_to_coco_api(val_loader.dataset, bbox_fmt='coco') all_models = glob.glob(os.path.join(Cfg.checkpoints, "*.pth")) all_models.sort(key=lambda fp: int( os.path.splitext(os.path.basename(fp))[0].replace("Yolov4_epoch", ""))) state_dict = torch.load(model_path, map_location=device) for model_path in all_models: print(f"eval on {os.path.splitext(os.path.basename(model_path))[0]}") model = Yolov4(None, 1, True) model.load_state_dict(state_dict) model.to(device) model.eval() coco_evaluator = CocoEvaluator(coco, iou_types=["bbox"]) for images, targets in val_loader: model_input = [[cv2.resize(img, (Cfg.w, Cfg.h))] for img in images] model_input = np.concatenate(model_input, axis=0) model_input = model_input.transpose(0, 3, 1, 2) model_input = torch.from_numpy(model_input).div(255.0) model_input = model_input.to(device) targets = [{k: v.to(device) for k, v in t.items()} for t in targets] if torch.cuda.is_available(): torch.cuda.synchronize() model_time = time.time() outputs = model(model_input) # outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs] model_time = time.time() - model_time outputs = outputs.cpu().detach().numpy() res = {} for img, target, output in zip(images, targets, outputs): img_height, img_width = img.shape[:2] boxes = output[..., :4].copy() # output boxes in yolo format boxes[..., :2] = boxes[ ..., :2] - boxes[..., 2:] / 2 # to coco format boxes[..., 0] = boxes[..., 0] * img_width boxes[..., 1] = boxes[..., 1] * img_height boxes[..., 2] = boxes[..., 2] * img_width boxes[..., 3] = boxes[..., 3] * img_height boxes = torch.as_tensor(boxes, dtype=torch.float32) labels = torch.as_tensor(np.zeros((len(output), )), dtype=torch.int64) scores = torch.as_tensor(output[..., -1], dtype=torch.float32) res[target["image_id"].item()] = { "boxes": boxes, "scores": scores, "labels": labels, } evaluator_time = time.time() coco_evaluator.update(res) evaluator_time = time.time() - evaluator_time # gather the stats from all processes coco_evaluator.synchronize_between_processes() # accumulate predictions from all images coco_evaluator.accumulate() coco_evaluator.summarize() del model del coco_evaluator
console.setLevel(log_level) formatter = logging.Formatter(fmt) console.setFormatter(formatter) logging.getLogger('').addHandler(console) return logging if __name__ == "__main__": logging = init_logger(log_dir='log') cfg = get_args(**Cfg) os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') model = Yolov4(cfg.pretrained) if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) model.to(device=device) try: train( model=model, config=cfg, epochs=cfg.TRAIN_EPOCHS, device=device, ) except KeyboardInterrupt: torch.save(model.state_dict(), 'INTERRUPTED.pth') logging.info('Saved interrupt')
help='path of your image file. or video or 0 to webcam', dest='source') args = parser.parse_args() return args if __name__ == '__main__': args = get_args() # model selection, darknet or pytorch if args.weightfile.split('.')[-1] == 'weights': # darknet m = Darknet(args.cfgfile) m.load_weights(args.weightfile) elif args.weightfile.split('.')[-1] in ['pth', 'pt', 'ph']: # pytorch m = Yolov4(n_classes=args.nclasses, inference=True) state_dict = torch.load(args.weightfile) m.load_state_dict(state_dict) m.head.inference = True m.eval() if use_cuda: torch.cuda.set_device(torch.device('cuda:1')) m.cuda() if args.source == '0': # webcam pass else: form = args.source.split('.')[-1] if form == 'jpg': # image pass
# import config file and save it to log update_config(cfg, args) init_seed(cfg.SEED) anchors = get_anchors(cfg) log_dir = os.path.join("log", os.path.basename(args.config_file)[:-5]) logging = init_logger(log_dir=log_dir) logging.info(pprint.pformat(args)) logging.info(cfg) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') if cfg.use_darknet_cfg: model = Darknet(cfg.cfgfile) else: model = Yolov4(anchors, yolov4conv137weight=cfg.pretrained, n_classes=cfg.classes) if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) model.to(device=device) try: train(model=model, config=cfg, tb_log_dir=log_dir, anchors=anchors, epochs=cfg.TRAIN_EPOCHS, device=device) except KeyboardInterrupt: torch.save(model.state_dict(), 'INTERRUPTED.pth')
def train(model, device, config, epochs=5, batch_size=1, save_cp=True, log_step=20, img_scale=0.5): train_dataset = Yolo_dataset(config.train_label, config, train=True) val_dataset = Yolo_dataset(config.val_label, config, train=False) n_train = len(train_dataset) n_val = len(val_dataset) train_loader = DataLoader(train_dataset, batch_size=config.batch // config.subdivisions, shuffle=True, num_workers=8, pin_memory=True, drop_last=True, collate_fn=collate) val_loader = DataLoader(val_dataset, batch_size=config.batch // config.subdivisions, shuffle=True, num_workers=8, pin_memory=True, drop_last=True, collate_fn=val_collate) writer = SummaryWriter( log_dir=config.TRAIN_TENSORBOARD_DIR, filename_suffix= f'OPT_{config.TRAIN_OPTIMIZER}_LR_{config.learning_rate}_BS_{config.batch}_Sub_{config.subdivisions}_Size_{config.width}', comment= f'OPT_{config.TRAIN_OPTIMIZER}_LR_{config.learning_rate}_BS_{config.batch}_Sub_{config.subdivisions}_Size_{config.width}' ) # writer.add_images('legend', # torch.from_numpy(train_dataset.label2colorlegend2(cfg.DATA_CLASSES).transpose([2, 0, 1])).to( # device).unsqueeze(0)) max_itr = config.TRAIN_EPOCHS * n_train # global_step = cfg.TRAIN_MINEPOCH * n_train global_step = 0 logging.info(f'''Starting training: Epochs: {epochs} Batch size: {config.batch} Subdivisions: {config.subdivisions} Learning rate: {config.learning_rate} Training size: {n_train} Validation size: {n_val} Checkpoints: {save_cp} Device: {device.type} Images size: {config.width} Optimizer: {config.TRAIN_OPTIMIZER} Dataset classes: {config.classes} Train label path:{config.train_label} Pretrained: ''') # learning rate setup def burnin_schedule(i): if i < config.burn_in: factor = pow(i / config.burn_in, 4) elif i < config.steps[0]: factor = 1.0 elif i < config.steps[1]: factor = 0.1 else: factor = 0.01 return factor if config.TRAIN_OPTIMIZER.lower() == 'adam': optimizer = optim.Adam( model.parameters(), lr=config.learning_rate / config.batch, betas=(0.9, 0.999), eps=1e-08, ) elif config.TRAIN_OPTIMIZER.lower() == 'sgd': optimizer = optim.SGD( params=model.parameters(), lr=config.learning_rate / config.batch, momentum=config.momentum, weight_decay=config.decay, ) scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule) criterion = Yolo_loss(device=device, batch=config.batch // config.subdivisions, n_classes=config.classes) # scheduler = ReduceLROnPlateau(optimizer, mode='max', verbose=True, patience=6, min_lr=1e-7) # scheduler = CosineAnnealingWarmRestarts(optimizer, 0.001, 1e-6, 20) save_prefix = 'Yolov4_epoch' saved_models = deque() model.train() for epoch in range(epochs): # model.train() epoch_loss = 0 epoch_step = 0 with tqdm(total=n_train, desc=f'Epoch {epoch + 1}/{epochs}', unit='img', ncols=50) as pbar: for i, batch in enumerate(train_loader): global_step += 1 epoch_step += 1 images = batch[0] bboxes = batch[1] images = images.to(device=device, dtype=torch.float32) bboxes = bboxes.to(device=device) bboxes_pred = model(images) loss, loss_xy, loss_wh, loss_obj, loss_cls, loss_l2 = criterion( bboxes_pred, bboxes) # loss = loss / config.subdivisions loss.backward() epoch_loss += loss.item() if global_step % config.subdivisions == 0: optimizer.step() scheduler.step() model.zero_grad() if global_step % (log_step * config.subdivisions) == 0: writer.add_scalar('train/Loss', loss.item(), global_step) writer.add_scalar('train/loss_xy', loss_xy.item(), global_step) writer.add_scalar('train/loss_wh', loss_wh.item(), global_step) writer.add_scalar('train/loss_obj', loss_obj.item(), global_step) writer.add_scalar('train/loss_cls', loss_cls.item(), global_step) writer.add_scalar('train/loss_l2', loss_l2.item(), global_step) writer.add_scalar('lr', scheduler.get_lr()[0] * config.batch, global_step) pbar.set_postfix( **{ 'loss (batch)': loss.item(), 'loss_xy': loss_xy.item(), 'loss_wh': loss_wh.item(), 'loss_obj': loss_obj.item(), 'loss_cls': loss_cls.item(), 'loss_l2': loss_l2.item(), 'lr': scheduler.get_lr()[0] * config.batch }) logging.debug( 'Train step_{}: loss : {},loss xy : {},loss wh : {},' 'loss obj : {},loss cls : {},loss l2 : {},lr : {}'. format(global_step, loss.item(), loss_xy.item(), loss_wh.item(), loss_obj.item(), loss_cls.item(), loss_l2.item(), scheduler.get_lr()[0] * config.batch)) pbar.update(images.shape[0]) if cfg.use_darknet_cfg: eval_model = Darknet(cfg.cfgfile, inference=True) else: eval_model = Yolov4(cfg.pretrained, n_classes=cfg.classes, inference=True) # eval_model = Yolov4(yolov4conv137weight=None, n_classes=config.classes, inference=True) if torch.cuda.device_count() > 1: eval_model.load_state_dict(model.module.state_dict()) else: eval_model.load_state_dict(model.state_dict()) eval_model.to(device) evaluator = evaluate(eval_model, val_loader, config, device) del eval_model stats = evaluator.coco_eval['bbox'].stats writer.add_scalar('train/AP', stats[0], global_step) writer.add_scalar('train/AP50', stats[1], global_step) writer.add_scalar('train/AP75', stats[2], global_step) writer.add_scalar('train/AP_small', stats[3], global_step) writer.add_scalar('train/AP_medium', stats[4], global_step) writer.add_scalar('train/AP_large', stats[5], global_step) writer.add_scalar('train/AR1', stats[6], global_step) writer.add_scalar('train/AR10', stats[7], global_step) writer.add_scalar('train/AR100', stats[8], global_step) writer.add_scalar('train/AR_small', stats[9], global_step) writer.add_scalar('train/AR_medium', stats[10], global_step) writer.add_scalar('train/AR_large', stats[11], global_step) if save_cp: try: # os.mkdir(config.checkpoints) os.makedirs(config.checkpoints, exist_ok=True) logging.info('Created checkpoint directory') except OSError: pass save_path = os.path.join(config.checkpoints, f'{save_prefix}{epoch + 1}.pth') torch.save(model.state_dict(), save_path) logging.info(f'Checkpoint {epoch + 1} saved !') saved_models.append(save_path) if len(saved_models) > config.keep_checkpoint_max > 0: model_to_remove = saved_models.popleft() try: os.remove(model_to_remove) except: logging.info(f'failed to remove {model_to_remove}') writer.close()
console.setLevel(log_level) formatter = logging.Formatter(fmt) console.setFormatter(formatter) logging.getLogger('').addHandler(console) return logging if __name__ == "__main__": logging = init_logger(log_dir='log') cfg = get_args(**Cfg) os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') model = Yolov4() if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) # model = model.cuda() model.to(device=device) try: train( model=model, config=cfg, epochs=cfg.TRAIN_EPOCHS, device=device, ) except KeyboardInterrupt: torch.save(model.state_dict(), 'INTERRUPTED.pth')