def detect(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: logging.warning("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # Load tested img imgfile = config["img_path"] img = Image.open(imgfile).convert('RGB') resized = img.resize((config["img_w"], config["img_h"])) input = image2torch(resized) input = input.to(torch.device("cuda")) start = time.time() outputs = net(input) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) output = non_max_suppression(output, config["yolo"]["classes"], conf_thres=0.5, nms_thres=0.4) finish = time.time() print('%s: Predicted in %f seconds.' % (imgfile, (finish - start))) namefile = config["classname_path"] class_names = load_class_names(namefile) plot_boxes(img, output, 'predictions.jpg', class_names)
def initial_yolo_model(config, size): is_training = False config["img_w"] = size config["img_h"] = size # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) if torch.cuda.is_available(): net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: if gpu: state_dict = torch.load(config["pretrain_snapshot"]) else: state_dict = torch.load(config["pretrain_snapshot"], map_location=torch.device('cpu')) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") return net
def test(config): is_training = False anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("load checkpoint from {}".format(config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append(YOLOLayer(config["batch_size"],i,config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # prepare images path images_name = os.listdir(config["images_path"]) images_path = [os.path.join(config["images_path"], name) for name in images_name] if len(images_path) == 0: raise Exception("no image found in {}".format(config["images_path"])) cap = cv2.VideoCapture(0) # cap = cv2.VideoCapture("./007.avi") img_i = 0 start = time.time() while cap.isOpened(): ret, frame = cap.read() if not ret: break img_i += 1 # preprocess images = [] images_origin = [] image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) images_origin.append(image) # keep for save result image = cv2.resize(image, (config["img_w"], config["img_h"]), interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) images = np.asarray(images) images = torch.from_numpy(images).cuda() # inference with torch.no_grad(): time1=datetime.datetime.now() outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) print("time1",(datetime.datetime.now()-time1).microseconds) batch_detections = non_max_suppression(output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"]) print("time2", (datetime.datetime.now() - time1).microseconds) # write result images. Draw bounding boxes and labels of detections classes = open(config["classes_names_path"], "r").read().split("\n")[:-1] if not os.path.isdir("./output/"): os.makedirs("./output/") for idx, detections in enumerate(batch_detections): img_show = images_origin[idx] img_show = cv2.cvtColor(img_show, cv2.COLOR_RGB2BGR) if detections is not None: unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) boxes=[] for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: # Rescale coordinates to original dimensions ori_h, ori_w = images_origin[idx].shape[:2] pre_h, pre_w = config["img_h"], config["img_w"] box_h = ((y2 - y1) / pre_h) * ori_h box_w = ((x2 - x1) / pre_w) * ori_w y1 = (y1 / pre_h) * ori_h x1 = (x1 / pre_w) * ori_w # Create a Rectangle patch box = BoundBox(x1, y1, x1 + box_w, y1 + box_h, cls_conf.item(), int(cls_pred)) boxes.append(box) img_show = draw_boxes(img_show, boxes, labels) # image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 255, 0), 1) cv2.imshow('1', img_show) cv2.waitKey(1) logging.info("Save all results to ./output/")
def test(config,int_dir='result'): is_training = False anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() ini_files = os.listdir(os.path.join(config['test_weights'], int_dir)) for kkk,ini_file in enumerate(ini_files): ini_list_config = configparser.ConfigParser() config_file_path = os.path.join(config['test_weights'], int_dir,ini_files[-kkk-1]) ini_list_config.read(config_file_path) ini_session = ini_list_config.sections() # accuracy = ini_list_config.items(ini_session[0]) err_jpgfiles = ini_list_config.items(ini_session[1]) aaa = glob.glob(os.path.join(config['test_weights'],'*_%s.weights'%ini_files[-kkk-1].split('_')[-1].split('.')[0])) weight_file = aaa[0]#os.path.join(config['test_weights'],'%s.weights'%ini_files[-kkk-1].split('_')[0]) if weight_file: # Restore pretrain model logging.info("load checkpoint from {}".format(weight_file)) state_dict = torch.load(weight_file) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") yolo_losses = [] for i in range(3): yolo_losses.append(YOLOLayer(1, i, config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) for index, _jpg_images in enumerate(err_jpgfiles): images = []# preprocess images_origin = [] jpg_path = str(_jpg_images[1]) print(str(index+1),jpg_path) bbox_list = read_gt_boxes(jpg_path) image = cv2.imread(jpg_path, cv2.IMREAD_COLOR) if image is None: logging.error("read path error: {}. skip it.".format(jpg_path)) continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) images_origin.append(image) # keep for save result image = cv2.resize(image, (config["img_w"], config["img_h"]),interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) images = np.asarray(images) images = torch.from_numpy(images).cuda() with torch.no_grad():# inference outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression(output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"]) classes = open(config["classes_names_path"], "r").read().split("\n")[:-1] if not os.path.isdir("./output/"): os.makedirs("./output/") for idx, detections in enumerate(batch_detections): image_show=images_origin[idx] if detections is not None: unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: ori_h, ori_w = images_origin[idx].shape[:2]# Rescale coordinates to original dimensions pre_h, pre_w = config["img_h"], config["img_w"] box_h = ((y2 - y1) / pre_h) * ori_h box_w = ((x2 - x1) / pre_w) * ori_w y1 = (y1 / pre_h) * ori_h x1 = (x1 / pre_w) * ori_w #绿色代表预测,红色代表标注 image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 255, 0),2) for (x1, x2, y1, y2) in bbox_list: [x1, x2, y1, y2] = map(int, [x1, x2, y1, y2]) cv2.rectangle(image_show, (x1, y1), (x2, y2), (0, 0, 255), 2) cv2.imshow('1', image_show) cv2.waitKey()
def test(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("load checkpoint from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # prepare images path images_name = os.listdir(config["images_path"]) images_path = [ os.path.join(config["images_path"], name) for name in images_name ] if len(images_path) == 0: raise Exception("no image found in {}".format(config["images_path"])) # Start inference batch_size = config["batch_size"] for step in range(0, len(images_path), batch_size): # preprocess images = [] images_origin = [] for path in images_path[step * batch_size:(step + 1) * batch_size]: logging.info("processing: {}".format(path)) image = cv2.imread(path, cv2.IMREAD_COLOR) if image is None: logging.error("read path error: {}. skip it.".format(path)) continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) images_origin.append(image) # keep for save result image = cv2.resize(image, (config["img_w"], config["img_h"]), interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) images = np.asarray(images) images = torch.from_numpy(images).cuda() # inference with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression( output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"]) # write result images. Draw bounding boxes and labels of detections classes = open(config["classes_names_path"], "r").read().split("\n")[:-1] if not os.path.isdir("./output/"): os.makedirs("./output/") for idx, detections in enumerate(batch_detections): plt.figure() fig, ax = plt.subplots(1) ax.imshow(images_origin[idx]) if detections is not None: unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) bbox_colors = random.sample(colors, n_cls_preds) for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: color = bbox_colors[int( np.where(unique_labels == int(cls_pred))[0])] # Rescale coordinates to original dimensions ori_h, ori_w = images_origin[idx].shape[:2] pre_h, pre_w = config["img_h"], config["img_w"] box_h = ((y2 - y1) / pre_h) * ori_h box_w = ((x2 - x1) / pre_w) * ori_w y1 = (y1 / pre_h) * ori_h x1 = (x1 / pre_w) * ori_w # Create a Rectangle patch bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolor='none') # Add the bbox to the plot ax.add_patch(bbox) # Add label plt.text(x1, y1, s=classes[int(cls_pred)], color='white', verticalalignment='top', bbox={ 'color': color, 'pad': 0 }) # Save generated image with detections plt.axis('off') plt.gca().xaxis.set_major_locator(NullLocator()) plt.gca().yaxis.set_major_locator(NullLocator()) plt.savefig('output/{}_{}.jpg'.format(step, idx), bbox_inches='tight', pad_inches=0.0) plt.close() logging.info("Save all results to ./output/")
def train(config): config["global_step"] = config.get("start_step", 0) is_training = False if config.get("export_onnx") else True anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Optimizer and learning rate optimizer = _get_optimizer(config, net) lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10) # lr_scheduler = optim.lr_scheduler.StepLR( # optimizer, # step_size=config["lr"]["decay_step"], # gamma=config["lr"]["decay_gamma"]) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load pretrained weights from {}".format(config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append(YOLOLayer(config["batch_size"],i,config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader dataloader = torch.utils.data.DataLoader(COCODataset(config["train_path"], (config["img_w"], config["img_h"]), is_training=True,is_scene=True), batch_size=config["batch_size"], shuffle=True,drop_last=True, num_workers=0, pin_memory=True) # Start the training loop logging.info("Start training.") dataload_len=len(dataloader) best_acc=0.5 for epoch in range(config["epochs"]): recall = 0 mini_step = 0 for step, samples in enumerate(dataloader): images, labels = samples["image"], samples["label"] start_time = time.time() config["global_step"] += 1 # Forward and backward optimizer.zero_grad() outputs = net(images) losses_name = ["total_loss", "x", "y", "w", "h", "conf", "cls", "recall"] losses = [0] * len(losses_name) for i in range(3): _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): losses[j] += l # losses = [sum(l) for l in losses] loss = losses[0] loss.backward() optimizer.step() _loss = loss.item() # example_per_second = config["batch_size"] / duration lr = optimizer.param_groups[0]['lr'] strftime = datetime.datetime.now().strftime("%H:%M:%S") # if (losses[7] / 3 >= recall / (step + 1)):#mini_batchΪ0×ßÕâÀï recall += losses[7] / 3 print('%s [Epoch %d/%d,batch %03d/%d loss:x %.5f,y %.5f,w %.5f,h %.5f,conf %.5f,cls %.5f,total %.5f,rec %.3f,avrec %.3f %.3f]' % (strftime, epoch, config["epochs"], step, dataload_len, losses[1], losses[2], losses[3], losses[4], losses[5], losses[6], _loss, losses[7] / 3, recall / (step + 1), lr)) if recall / len(dataloader) > best_acc: best_acc=recall / len(dataloader) if epoch>0: torch.save(net.state_dict(), '%s/%.4f_%04d.weights' % (checkpoint_dir, recall / len(dataloader), epoch)) lr_scheduler.step() net.train(is_training) torch.cuda.empty_cache() # net.train(True) logging.info("Bye bye")
def main(video_fn): logging.basicConfig(level=logging.DEBUG, format="[%(asctime)s %(filename)s] %(message)s") if len(sys.argv) != 2: logging.error("Usage: python video.py params.py") sys.exit() params_path = sys.argv[1] if not os.path.isfile(params_path): logging.error("no params file found! path: {}".format(params_path)) sys.exit() config = importlib.import_module(params_path[:-3]).TRAINING_PARAMS config["batch_size"] *= len(config["parallels"]) is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # load pretrained model if config["pretrain_snapshot"]: logging.info("load checkpoint from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # load class names classes = open(config["classes_names_path"], "r").read().split("\n")[:-1] cap = cv2.VideoCapture(video_fn) # Check if camera opened successfully if (cap.isOpened() == False): print("Error opening video stream or file") # Read until video is completed while (cap.isOpened()): # Capture frame-by-frame ret, frame = cap.read() frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5) if ret == True: # 1. pre-process image logging.info("processing frame") image_tensor = prep_image(frame, config) with torch.no_grad(): outputs = net(image_tensor) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression( output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"], nms_thres=0.45) for idx, detections in enumerate(batch_detections): if detections is not None: unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) bbox_colors = random.sample(colors, n_cls_preds) for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: color = bbox_colors[int( np.where(unique_labels == int(cls_pred))[0])] # Rescale coordinates to original dimensions x1, y1, box_w, box_h = get_rescaled_coords( frame.shape[0], frame.shape[1], config["img_h"], config["img_w"], x1, y1, x2, y2) cv2.rectangle(frame, (x1, y1), (x1 + box_w, y1 + box_h), color, 2) cv2.putText(frame, classes[int(cls_pred)], (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 1, cv2.LINE_AA) cv2.imshow('Frame', frame) # Press Q on keyboard to exit if cv2.waitKey(25) & 0xFF == ord('q'): break # Break the loop else: break # When everything done, release the video capture object cap.release() # Closes all the frames cv2.destroyAllWindows()
def train(config): config["global_step"] = config.get("start_step", 0) is_training = False if config.get("export_onnx") else True # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Optimizer and learning rate optimizer = _get_optimizer(config, net) lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=15) # lr_scheduler = optim.lr_scheduler.StepLR( # optimizer, # step_size=config["lr"]["decay_step"], # gamma=config["lr"]["decay_gamma"]) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load pretrained weights from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) # Start the training loop logging.info("Start training.") dataload_len = len(dataloader) epoch_size = 4 start = time.time() pruned_pct = 0 global index, pruned_book, num_pruned global num_weights global weight_masks, bias_masks for epoch in range(config["epochs"]): if epoch % 4 == 0: index = 0 num_pruned = 0 num_weights = 0 net.apply(prune) torch.save(net.state_dict(), '%s/%.4f_%04d.weights' % (checkpoint_dir, 0.01, 1)) print('previously pruned: %.3f %%' % (100 * (pruned_pct))) print('number pruned: %.3f %%' % (100 * (num_pruned / num_weights))) new_pruned = num_pruned / num_weights - pruned_pct pruned_pct = num_pruned / num_weights # if new_pruned <= 0.01: # time_elapse = time.time() - start # print('training time:', str(timedelta(seconds=time_elapse))) # break recall = 0 mini_step = 0 for step, samples in enumerate(dataloader): index = 0 images, labels = samples["image"], samples["label"] start_time = time.time() optimizer.zero_grad() outputs = net(images) losses_name = [ "total_loss", "x", "y", "w", "h", "conf", "cls", "recall" ] losses = [0] * len(losses_name) for i in range(3): _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): losses[j] += l # losses = [sum(l) for l in losses] loss = losses[0] loss.backward() net.apply(set_grad) optimizer.step() _loss = loss.item() # example_per_second = config["batch_size"] / duration lr = optimizer.param_groups[0]['lr'] strftime = datetime.datetime.now().strftime("%H:%M:%S") recall += losses[7] / 3 print( '%s [Epoch %d/%d,batch %03d/%d loss:x %.5f,y %.5f,w %.5f,h %.5f,conf %.5f,cls %.5f,total %.5f,rec %.3f,avrec %.3f %.3f]' % (strftime, epoch, config["epochs"], step, dataload_len, losses[1], losses[2], losses[3], losses[4], losses[5], losses[6], _loss, losses[7] / 3, recall / (step + 1), lr)) if (epoch % 2 == 0 and recall / len(dataloader) > 0.5 ) or recall / len(dataloader) > 0: # torch.save(net.state_dict(), '%s/%.4f_%04d.weights' % (checkpoint_dir, recall / len(dataloader), epoch)) torch.save( net.state_dict(), '%s/%.4f_%04d.weights' % (checkpoint_dir, recall / len(dataloader), epoch)) lr_scheduler.step() # net.train(True) logging.info("Bye bye")
def evaluate(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load checkpoint: {}".format(config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: logging.warning("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader. dataloader = torch.utils.data.DataLoader(COCODataset( config["val_path"], (config["img_w"], config["img_h"]), is_training=False), batch_size=config["batch_size"], shuffle=False, num_workers=8, pin_memory=False) # Coco Prepare. index2category = json.load(open("coco_index2category.json")) # Start the eval loop logging.info("Start eval.") coco_results = [] coco_img_ids = set([]) APs = [] for step, samples in enumerate(dataloader): images, labels = samples["image"], samples["label"] image_paths, origin_sizes = samples["image_path"], samples[ "origin_size"] with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression(output, config["yolo"]["classes"], conf_thres=0.0001, nms_thres=0.45) for idx, detections in enumerate(batch_detections): correct = [] annotations = labels[idx, labels[idx, :, 3] != 0] image_id = int(os.path.basename(image_paths[idx])[-16:-4]) coco_img_ids.add(image_id) if detections is None: if annotations.size(0) != 0: APs.append(0) continue detections = detections[np.argsort(-detections[:, 4])] origin_size = eval(origin_sizes[idx]) detections = detections.cpu().numpy() # =========================================================================================================================== # The amount of padding that was added pad_x = max(origin_size[1] - origin_size[0], 0) * (config["img_w"] / max(origin_size)) pad_y = max(origin_size[0] - origin_size[1], 0) * (config["img_w"] / max(origin_size)) # Image height and width after padding is removed unpad_h = config["img_w"] - pad_y unpad_w = config["img_w"] - pad_x # =========================================================================================================================== if annotations.size(0) == 0: correct.extend([0 for _ in range(len(detections))]) else: target_boxes = torch.FloatTensor(annotations[:, 1:].shape) target_boxes[:, 0] = (annotations[:, 1] - annotations[:, 3] / 2) target_boxes[:, 1] = (annotations[:, 2] - annotations[:, 4] / 2) target_boxes[:, 2] = (annotations[:, 1] + annotations[:, 3] / 2) target_boxes[:, 3] = (annotations[:, 2] + annotations[:, 4] / 2) target_boxes *= config["img_w"] detected = [] for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: pred_bbox = (x1, y1, x2, y2) #x1 = x1 / config["img_w"] * origin_size[0] #x2 = x2 / config["img_w"] * origin_size[0] #y1 = y1 / config["img_h"] * origin_size[1] #y2 = y2 / config["img_h"] * origin_size[1] #w = x2 - x1 #h = y2 - y1 h = ((y2 - y1) / unpad_h) * origin_size[1] w = ((x2 - x1) / unpad_w) * origin_size[0] y1 = ((y1 - pad_y // 2) / unpad_h) * origin_size[1] x1 = ((x1 - pad_x // 2) / unpad_w) * origin_size[0] coco_results.append({ "image_id": image_id, "category_id": index2category[str(int(cls_pred.item()))], "bbox": (float(x1), float(y1), float(w), float(h)), "score": float(conf), }) pred_bbox = torch.FloatTensor(pred_bbox).view(1, -1) # Compute iou with target boxes iou = bbox_iou(pred_bbox, target_boxes) # Extract index of largest overlap best_i = np.argmax(iou) # If overlap exceeds threshold and classification is correct mark as correct if iou[best_i] > config[ 'iou_thres'] and cls_pred == annotations[ best_i, 0] and best_i not in detected: correct.append(1) detected.append(best_i) else: correct.append(0) true_positives = np.array(correct) false_positives = 1 - true_positives # Compute cumulative false positives and true positives false_positives = np.cumsum(false_positives) true_positives = np.cumsum(true_positives) # Compute recall and precision at all ranks recall = true_positives / annotations.size(0) if annotations.size( 0) else true_positives precision = true_positives / np.maximum( true_positives + false_positives, np.finfo(np.float64).eps) # Compute average precision AP = compute_ap(recall, precision) APs.append(AP) print("+ Sample [%d/%d] AP: %.4f (%.4f)" % (len(APs), 5000, AP, np.mean(APs))) logging.info("Now {}/{}".format(step, len(dataloader))) print("Mean Average Precision: %.4f" % np.mean(APs)) save_results_path = "coco_results.json" with open(save_results_path, "w") as f: json.dump(coco_results, f, sort_keys=True, indent=4, separators=(',', ':')) logging.info("Save coco format results to {}".format(save_results_path)) # COCO api logging.info("Using coco-evaluate tools to evaluate.") cocoGt = COCO(config["annotation_path"]) cocoDt = cocoGt.loadRes(save_results_path) cocoEval = COCOeval(cocoGt, cocoDt, "bbox") cocoEval.params.imgIds = list(coco_img_ids) # real imgIds cocoEval.evaluate() cocoEval.accumulate() cocoEval.summarize()
def train(imgs, labels, checkpoint_path, config): config["global_step"] = config.get("start_step", 0) is_training = False if config.get("export_onnx") else True # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Optimizer and learning rate optimizer = _get_optimizer(config, net) lr_scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=config["lr"]["decay_step"], gamma=config["lr"]["decay_gamma"]) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if checkpoint_path: logging.info("Load pretrained weights from {}".format(checkpoint_path)) state_dict = torch.load(checkpoint_path) net.load_state_dict(state_dict) # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader dataloader = torch.utils.data.DataLoader(SatDataset( imgs, labels, (config["img_w"], config["img_h"]), is_training=True), batch_size=config["batch_size"], shuffle=True, num_workers=1, pin_memory=True) # Start the training loop logging.info("Start training.") for epoch in range(config["epochs"]): for step, samples in enumerate(dataloader): images, labels = samples["image"], samples["label"] start_time = time.time() config["global_step"] += 1 # Forward and backward optimizer.zero_grad() outputs = net(images) losses_name = ["total_loss", "x", "y", "w", "h", "conf", "cls"] losses = [[]] * len(losses_name) for i in range(3): _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): losses[j].append(l) losses = [sum(l) for l in losses] loss = losses[0] loss.backward() optimizer.step() if step > 0 and step % 10 == 0: _loss = loss.item() duration = float(time.time() - start_time) example_per_second = config["batch_size"] / duration lr = optimizer.param_groups[0]['lr'] logging.info( "epoch [%.3d] iter = %d loss = %.2f example/sec = %.3f lr = %.5f " % (epoch, step, _loss, example_per_second, lr)) config["tensorboard_writer"].add_scalar( "lr", lr, config["global_step"]) config["tensorboard_writer"].add_scalar( "example/sec", example_per_second, config["global_step"]) for i, name in enumerate(losses_name): value = _loss if i == 0 else losses[i] config["tensorboard_writer"].add_scalar( name, value, config["global_step"]) lr_scheduler.step() # net.train(False) checkpoint_path = _save_checkpoint(net.state_dict(), config) # net.train(True) logging.info("Bye~") return checkpoint_path
def test(config): is_training = False anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("load checkpoint from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # prepare images path images_name = os.listdir(config["images_path"]) images_path = [ os.path.join(config["images_path"], name) for name in images_name ] if len(images_path) == 0: raise Exception("no image found in {}".format(config["images_path"])) # Start inference batch_size = config["batch_size"] for step in range(0, len(images_path), batch_size): # preprocess images = [] images_origin = [] for path in images_path[step * batch_size:(step + 1) * batch_size]: logging.info("processing: {}".format(path)) image = cv2.imread(path, cv2.IMREAD_COLOR) if image is None: logging.error("read path error: {}. skip it.".format(path)) continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) images_origin.append(image) # keep for save result image = cv2.resize(image, (config["img_w"], config["img_h"]), interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) images = np.asarray(images) images = torch.from_numpy(images).cuda() # inference with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression( output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"]) # write result images. Draw bounding boxes and labels of detections classes = open(config["classes_names_path"], "r").read().split("\n")[:-1] for idx, detections in enumerate(batch_detections): if detections is not None: unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) boxes = [] for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: # Rescale coordinates to original dimensions ori_h, ori_w = images_origin[idx].shape[:2] pre_h, pre_w = config["img_h"], config["img_w"] box_h = ((y2 - y1) / pre_h) * ori_h box_w = ((x2 - x1) / pre_w) * ori_w y1 = (y1 / pre_h) * ori_h x1 = (x1 / pre_w) * ori_w # Create a Rectangle patch box = BoundBox(x1, y1, x1 + box_w, y1 + box_h, cls_conf.item(), int(cls_pred), classes[int(cls_pred)]) boxes.append(box) # Save generated image with detections img_show = draw_boxes(images_origin[idx], boxes, labels, 0.5) img_show = cv2.resize(img_show, (img_show.shape[1], img_show.shape[0]), interpolation=cv2.INTER_CUBIC) # outVideo.write(img_show) cv2.imshow("ai", img_show) cv2.waitKey() logging.info("Save all results to ./output/")
def train(config): config["global_step"] = config.get("start_step", 0) is_training = False if config.get("export_onnx") else True anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Optimizer and learning rate optimizer = _get_optimizer(config, net) t_max = 50 # lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=t_max,eta_min=1e-05) lr_scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=config["lr"]["decay_step"], gamma=config["lr"]["decay_gamma"]) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load pretrained weights from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) # Only export onnx # if config.get("export_onnx"): # real_model = net.module # real_model.eval() # dummy_input = torch.randn(8, 3, config["img_h"], config["img_w"]).cuda() # save_path = os.path.join(config["sub_working_dir"], "pytorch.onnx") # logging.info("Exporting onnx to {}".format(save_path)) # torch.onnx.export(real_model, dummy_input, save_path, verbose=False) # logging.info("Done. Exiting now.") # sys.exit() # Evaluate interface # if config["evaluate_type"]: # logging.info("Using {} to evaluate model.".format(config["evaluate_type"])) # evaluate_func = importlib.import_module(config["evaluate_type"]).run_eval # config["online_net"] = net # YOLO loss with 3 scales # DataLoader dataloader = torch.utils.data.DataLoader( COCODataset(config["train_path"], (config["img_w"], config["img_h"]), is_training=True, is_scene=True), batch_size=config["batch_size"] * config["parallels"], shuffle=True, drop_last=True, num_workers=0, pin_memory=True) # Start the training loop logging.info("Start training.") dataload_len = len(dataloader) best_acc = 0.2 last_recall = 0.6 for epoch in range(config["epochs"]): recall = 0 mini_step = 0 for step, samples in enumerate(dataloader): start = time.time() images, labels = samples["image"], samples["label"] config["global_step"] += 1 # Forward and backward optimizer.zero_grad() losses = net(images.cuda(), labels.cuda()) # current_recall = mAP(detections, labels, config["img_w"]) # current_recall = np.mean(current_recall) if config["parallels"] > 1: losses = losses.view(config["parallels"], 8)[0] + losses.view( config["parallels"], 8)[1] loss = losses[0] if epoch > 0: loss = loss * 20 current_recall = float(losses[7] / 3 / config["parallels"]) if last_recall < 0.65: loss = loss + 20 * (1 - current_recall) # * 0.8 else: loss = loss + 20 * (1 - current_recall) loss.backward() optimizer.step() _loss = loss.item() # example_per_second = config["batch_size"] / duration lr = optimizer.param_groups[0]['lr'] # strftime = datetime.datetime.now().strftime("%H:%M:%S") # # if (losses[7] / 3 >= recall / (step + 1)):#mini_batch为0走这里 recall += current_recall print( '%s [Epoch %d/%d,batch %03d/%d loss:x %.5f,y %.5f,w %.5f,h %.5f,conf %.5f,cls %.5f,total %.5f,rec %.3f,avrec %.3f %.3f]' % (strftime, epoch, config["epochs"], step, dataload_len, losses[1], losses[2], losses[3], losses[4], losses[5], losses[6], _loss, current_recall, recall / (step + 1), lr)) last_recall = recall / len(dataloader) if recall / len(dataloader) > best_acc: best_acc = recall / len(dataloader) torch.save( net.state_dict(), '%s/%.4f_%04d.weights' % (checkpoint_dir, recall / len(dataloader), epoch)) lr_scheduler.step() # if epoch % (lr_scheduler.T_max + next_need) == (lr_scheduler.T_max + next_need - 1): # next_need += float(lr_scheduler.T_max) # lr_scheduler.T_max += 2 # lr_scheduler.last_epoch = 0 # lr_scheduler.base_lrs*=0.98 # lr_scheduler.base_lrs[0] *= 0.95 # lr_scheduler.base_lrs[1] *= 0.95 # net.train(is_training) # torch.cuda.empty_cache() # net.train(True) logging.info("Bye bye")
def train(config): config["global_step"] = config.get("start_step", 0) is_training = False if config.get("export_onnx") else True anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Optimizer and learning rate optimizer = _get_optimizer(config, net) lr_scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=config["lr"]["decay_step"], gamma=config["lr"]["decay_gamma"]) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load pretrained weights from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) total_loss = 0 last_total_loss = 0 manager = Manager() # 父进程创建Queue,并传给各个子进程: q = manager.Queue(1) lock = manager.Lock() # 初始化一把锁 p = Pool() pw = p.apply_async(get_data, args=(q, lock)) batch_len = q.get() if batch_len[0] == "len": batch_len = batch_len[1] logging.info("Start training.") for epoch in range(config["epochs"]): recall = 0 for step in range(batch_len): samples = q.get() images, labels = samples["image"], samples["label"] start_time = time.time() config["global_step"] += 1 # Forward and backward optimizer.zero_grad() outputs = net(images) losses_name = [ "total_loss", "x", "y", "w", "h", "conf", "cls", "recall" ] losses = [0] * len(losses_name) for i in range(3): _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): losses[j] += l # losses = [sum(l) for l in losses] loss = losses[0] loss.backward() optimizer.step() if step > 0 and step % 2 == 0: _loss = loss.item() duration = float(time.time() - start_time) example_per_second = config["batch_size"] / duration lr = optimizer.param_groups[0]['lr'] strftime = datetime.datetime.now().strftime("%H:%M:%S") recall += losses[7] / 3 print( '%s [Epoch %d/%d, Batch %03d/%d losses: x %.5f, y %.5f, w %.5f, h %.5f, conf %.5f, cls %.5f, total %.5f, recall: %.3f]' % (strftime, epoch, config["epochs"], step, batch_len, losses[1], losses[2], losses[3], losses[4], losses[5], losses[6], _loss, losses[7] / 3)) # logging.info(epoch [%.3d] iter = %d loss = %.2f example/sec = %.3f lr = %.5f "% # (epoch, step, _loss, example_per_second, lr)) # config["tensorboard_writer"].add_scalar("lr", # lr, # config["global_step"]) # config["tensorboard_writer"].add_scalar("example/sec", # example_per_second, # config["global_step"]) # for i, name in enumerate(losses_name): # value = _loss if i == 0 else losses[i] # config["tensorboard_writer"].add_scalar(name, # value, # config["global_step"]) if (epoch % 2 == 0 and recall / batch_len > 0.7) or recall / batch_len > 0.96: torch.save(net.state_dict(), '%s/%04d.weights' % (checkpoint_dir, epoch)) lr_scheduler.step()
def test(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("load checkpoint from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLayer(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # prepare images path images_name = os.listdir(config["images_path"]) images_path = [ os.path.join(config["images_path"], name) for name in images_name ] if len(images_path) == 0: raise Exception("no image found in {}".format(config["images_path"])) # Start testing FPS of different batch size for batch_size in range(1, 10): # preprocess images = [] for path in images_path[:batch_size]: image = cv2.imread(path, cv2.IMREAD_COLOR) if image is None: logging.error("read path error: {}. skip it.".format(path)) continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = cv2.resize(image, (config["img_w"], config["img_h"]), interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) for i in range(batch_size - len(images)): images.append(images[0]) # fill len to batch_sze images = np.asarray(images) images = torch.from_numpy(images).cuda() # inference in 30 times and calculate average inference_times = [] for i in range(30): start_time = time.time() with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression( output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"]) torch.cuda.synchronize() # wait all done. end_time = time.time() inference_times.append(end_time - start_time) inference_time = sum(inference_times) / len( inference_times) / batch_size fps = 1.0 / inference_time logging.info( "Batch_Size: {}, Inference_Time: {:.5f} s/image, FPS: {}".format( batch_size, inference_time, fps))
def evaluate(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: logging.warning("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append(YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader dataloader = torch.utils.data.DataLoader(dataset=COCODataset(config["val_path"], config["img_w"]), batch_size=config["batch_size"], shuffle=True, num_workers=1, pin_memory=False) # Start the eval loop logging.info("Start eval.") n_gt = 0 correct = 0 logging.info('%s' % str(dataloader)) gt_histro={} pred_histro = {} correct_histro = {} for i in range(config["yolo"]["classes"]): gt_histro[i] = 1 pred_histro[i] = 1 correct_histro[i] = 0 # images 是一个batch里的全部图片,labels是一个batch里面的全部标签 for step, (images, labels) in enumerate(dataloader): labels = labels.cuda() with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) # 把三个尺度上的预测结果在第1维度(第0维度是batch里的照片,第1维度是一张照片里面的各个预测框,第2维度是各个预测数值)上拼接起来 batch_output = torch.cat(output_list, dim=1) logging.info('%s' % str(batch_output.shape)) # 进行NMS抑制 batch_output = non_max_suppression(prediction=batch_output, num_classes=config["yolo"]["classes"], conf_thres=config["conf_thresh"], nms_thres=config["nms_thresh"]) # calculate for sample_index_in_batch in range(labels.size(0)): # fetched img sample in tensor( C(RxGxB) x H x W ), transform to cv2 format in H x W x C(BxGxR) sample_image = images[sample_index_in_batch].numpy() sample_image = np.transpose(sample_image, (1, 2, 0)) sample_image = cv2.cvtColor(sample_image, cv2.COLOR_RGB2BGR) logging.debug("fetched img %d size %s" % (sample_index_in_batch, sample_image.shape)) # Get labels for sample where width is not zero (dummies)(init all labels to zeros in array) target_sample = labels[sample_index_in_batch, labels[sample_index_in_batch, :, 3] != 0] # get prediction for this sample sample_pred = batch_output[sample_index_in_batch] if sample_pred is not None: for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred: # for each prediction box # logging.info("%d" % obj_cls) box_pred = torch.cat([coord.unsqueeze(0) for coord in [x1, y1, x2, y2]]).view(1, -1) sample_image = draw_prediction(sample_image,conf, obj_conf, int(obj_pred), (x1, y1, x2, y2), config) # 每一个ground truth的 分类编号obj_cls、相对中心x、相对中心y、相对宽w、相对高h for obj_cls, tx, ty, tw, th in target_sample: # Get rescaled gt coordinates # 转化为输入像素尺寸的 左上角像素tx1 ty1,右下角像素tx2 ty2 tx1, tx2 = config["img_w"] * (tx - tw / 2), config["img_w"] * (tx + tw / 2) ty1, ty2 = config["img_h"] * (ty - th / 2), config["img_h"] * (ty + th / 2) # 计算ground truth数量,用于统计信息 n_gt += 1 gt_histro[int(obj_cls)] += 1 # 转化为 shape(1,4)的tensor,用来计算IoU box_gt = torch.cat([coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2]]).view(1, -1) # logging.info('%s' % str(box_gt.shape)) sample_pred = batch_output[sample_index_in_batch] if sample_pred is not None: # Iterate through predictions where the class predicted is same as gt # 对于每一个ground truth,遍历预测结果 for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[sample_pred[:, 6] == obj_cls]: # 如果当前预测分类 == 当前真实分类 #logging.info("%d" % obj_cls) box_pred = torch.cat([coord.unsqueeze(0) for coord in [x1, y1, x2, y2]]).view(1, -1) pred_histro[int(obj_pred)] += 1 iou = bbox_iou(box_pred, box_gt) if iou >= config["iou_thresh"]: correct += 1 correct_histro[int(obj_pred)] += 1 break if n_gt: types = config["types"] reverse_types = {} # 建立一个反向的types for key in types.keys(): reverse_types[types[key]] = key logging.info('Batch [%d/%d] mAP: %.5f' % (step, len(dataloader), float(correct / n_gt))) logging.info('mAP Histro:%s' % str([ reverse_types[i] +':'+ str(int(100 * correct_histro[i] / gt_histro[i])) for i in range(config["yolo"]["classes"] ) ])) logging.info('Recall His:%s' % str([ reverse_types[i] +':'+ str(int(100 * correct_histro[i] / pred_histro[i])) for i in range(config["yolo"]["classes"]) ])) logging.info('Mean Average Precision: %.5f' % float(correct / n_gt))
def train(config): config["global_step"] = config.get("start_step", 0) is_training = False if config.get("export_onnx") else True anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Optimizer and learning rate optimizer = _get_optimizer(config, net) lr_scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=config["lr"]["decay_step"], gamma=config["lr"]["decay_gamma"]) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load pretrained weights from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) # Only export onnx # if config.get("export_onnx"): # real_model = net.module # real_model.eval() # dummy_input = torch.randn(8, 3, config["img_h"], config["img_w"]).cuda() # save_path = os.path.join(config["sub_working_dir"], "pytorch.onnx") # logging.info("Exporting onnx to {}".format(save_path)) # torch.onnx.export(real_model, dummy_input, save_path, verbose=False) # logging.info("Done. Exiting now.") # sys.exit() # Evaluate interface # if config["evaluate_type"]: # logging.info("Using {} to evaluate model.".format(config["evaluate_type"])) # evaluate_func = importlib.import_module(config["evaluate_type"]).run_eval # config["online_net"] = net # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader dataloader = torch.utils.data.DataLoader(COCODataset( config["train_path"], (config["img_w"], config["img_h"]), is_training=True, is_scene=True), batch_size=config["batch_size"], shuffle=True, drop_last=True, num_workers=0, pin_memory=True) # Start the training loop logging.info("Start training.") dataload_len = len(dataloader) for epoch in range(config["epochs"]): recall = 0 mini_step = 0 for step, samples in enumerate(dataloader): images, labels = samples["image"], samples["label"] start_time = time.time() config["global_step"] += 1 for mini_batch in range(3): mini_step += 1 # Forward and backward optimizer.zero_grad() outputs = net(images) losses_name = [ "total_loss", "x", "y", "w", "h", "conf", "cls", "recall" ] losses = [0] * len(losses_name) for i in range(3): _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): losses[j] += l # losses = [sum(l) for l in losses] loss = losses[0] loss.backward() optimizer.step() _loss = loss.item() # example_per_second = config["batch_size"] / duration # lr = optimizer.param_groups[0]['lr'] strftime = datetime.datetime.now().strftime("%H:%M:%S") if (losses[7] / 3 >= recall / (step + 1)) or mini_batch == (3 - 1): #mini_batch为0走这里 recall += losses[7] / 3 print( '%s [Epoch %d/%d,batch %03d/%d loss:x %.5f,y %.5f,w %.5f,h %.5f,conf %.5f,cls %.5f,total %.5f,rec %.3f,avrec %.3f %d]' % (strftime, epoch, config["epochs"], step, dataload_len, losses[1], losses[2], losses[3], losses[4], losses[5], losses[6], _loss, losses[7] / 3, recall / (step + 1), mini_batch)) break else: print( '%s [Epoch %d/%d,batch %03d/%d loss:x %.5f,y %.5f,w %.5f,h %.5f,conf %.5f,cls %.5f,total %.5f,rec %.3f,prerc %.3f %d]' % (strftime, epoch, config["epochs"], step, dataload_len, losses[1], losses[2], losses[3], losses[4], losses[5], losses[6], _loss, losses[7] / 3, recall / step, mini_batch)) # logging.info(epoch [%.3d] iter = %d loss = %.2f example/sec = %.3f lr = %.5f "% # (epoch, step, _loss, example_per_second, lr)) # config["tensorboard_writer"].add_scalar("lr", # lr, # config["global_step"]) # config["tensorboard_writer"].add_scalar("example/sec", # example_per_second, # config["global_step"]) # for i, name in enumerate(losses_name): # value = _loss if i == 0 else losses[i] # config["tensorboard_writer"].add_scalar(name, # value, # config["global_step"]) if (epoch % 2 == 0 and recall / len(dataloader) > 0.7 ) or recall / len(dataloader) > 0.96: torch.save( net.state_dict(), '%s/%.4f_%04d.weights' % (checkpoint_dir, recall / len(dataloader), epoch)) lr_scheduler.step() # net.train(True) logging.info("Bye bye")
def train(): global_step = 0 is_training = True # Load and Initialize Network net = ModelMain(is_training) net.train(is_training) # Optimizer and Lr optimizer = _get_optimizer(net) lr_scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=lr_decay_step, #20 gamma=lr_decay_gamma) # 0.1 # Set Data Paraller: net = nn.DataParallel(net) net = net.cuda() logging.info("Net of Cuda is Done!") # Restore pretrain model 从预训练模型中恢复 if pretrain_snapshot: logging.info( "Load pretrained weights from {}".format(pretrain_snapshot)) state_dic = torch.load(pretrain_snapshot) net.load_state_dict(state_dic) yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(anchors[i], classes, (img_w, img_h)).cuda()) print('YOLO_Losses: \n', yolo_losses) # DataLoader train_data_loader = DATA.DataLoader(dataset=COCODataset(train_path, (img_w, img_h), is_training=True), batch_size=batch_size, shuffle=True, pin_memory=False) # Start the training loop logging.info("Start training......") for epoch in range(epochs): for step, samples in enumerate(train_data_loader): images, labels = samples['image'].cuda(), samples["label"].cuda() start_time = time.time() global_step += 1 # Forward & Backward optimizer.zero_grad() outputs = net(images) losses_name = ["total_loss", "x", "y", "w", "h", "conf", "cls"] losses = [[]] * len( losses_name) # [[]] ---> [[], [], [], [], [], [], []] for i in range(3): # YOLO 3 scales _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): # print('j: ', j, 'l: ', l) j: index(0-6); l内容: 总loss, x, y, w, h, conf, cls losses[j].append(l) losses = [sum(l) for l in losses] loss = losses[0] # losses[0]为总Loss conf = losses[5] loss.backward() optimizer.step() if step > 0 and step % 10 == 0: _loss = loss.item() _conf = conf.item() duration = float(time.time() - start_time) # 总用时 example_per_second = batch_size / duration # 每个样本用时 lr = optimizer.param_groups[0]['lr'] logging.info( "epoch [%.3d] iter = %d loss = %.2f conf = %.2f example/sec = %.3f lr = %.5f " % (epoch, step, _loss, _conf, example_per_second, lr)) if step >= 0 and step % 1000 == 0: # net.train(False) _save_checkpoint(net.state_dict(), epoch, step) # net.train(True) lr_scheduler.step() _save_checkpoint(net.state_dict(), 100, 9999) logging.info("Bye~")
def evaluate(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: logging.warning("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader dataloader = torch.utils.data.DataLoader(COCODataset( config["val_path"], (config["img_w"], config["img_h"]), is_training=False), batch_size=config["batch_size"], shuffle=False, num_workers=16, pin_memory=False) # Start the eval loop logging.info("Start eval.") n_gt = 0 correct = 0 for step, samples in enumerate(dataloader): images, labels = samples["image"], samples["label"] labels = labels.cuda() with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) output = non_max_suppression(output, 80, conf_thres=0.2) # calculate for sample_i in range(labels.size(0)): # Get labels for sample where width is not zero (dummies) target_sample = labels[sample_i, labels[sample_i, :, 3] != 0] for obj_cls, tx, ty, tw, th in target_sample: # Get rescaled gt coordinates tx1, tx2 = config["img_w"] * ( tx - tw / 2), config["img_w"] * (tx + tw / 2) ty1, ty2 = config["img_h"] * ( ty - th / 2), config["img_h"] * (ty + th / 2) n_gt += 1 box_gt = torch.cat([ coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2] ]).view(1, -1) sample_pred = output[sample_i] if sample_pred is not None: # Iterate through predictions where the class predicted is same as gt for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[ sample_pred[:, 6] == obj_cls]: box_pred = torch.cat([ coord.unsqueeze(0) for coord in [x1, y1, x2, y2] ]).view(1, -1) iou = bbox_iou(box_pred, box_gt) if iou >= config["iou_thres"]: correct += 1 break if n_gt: logging.info('Batch [%d/%d] mAP: %.5f' % (step, len(dataloader), float(correct / n_gt))) logging.info('Mean Average Precision: %.5f' % float(correct / n_gt))
def evaluate(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load checkpoint: {}".format(config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: logging.warning("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader. dataloader = torch.utils.data.DataLoader(COCODataset( config["val_path"], (config["img_w"], config["img_h"]), is_training=False), batch_size=config["batch_size"], shuffle=False, num_workers=8, pin_memory=False) # Coco Prepare. index2category = json.load(open("coco_index2category.json")) # Start the eval loop logging.info("Start eval.") coco_results = [] coco_img_ids = set([]) for step, samples in enumerate(dataloader): images, labels = samples["image"], samples["label"] image_paths, origin_sizes = samples["image_path"], samples[ "origin_size"] with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression(output, config["yolo"]["classes"], conf_thres=0.01, nms_thres=0.45) for idx, detections in enumerate(batch_detections): image_id = int(os.path.basename(image_paths[idx])[-16:-4]) coco_img_ids.add(image_id) if detections is not None: origin_size = eval(origin_sizes[idx]) detections = detections.cpu().numpy() for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: x1 = x1 / config["img_w"] * origin_size[0] x2 = x2 / config["img_w"] * origin_size[0] y1 = y1 / config["img_h"] * origin_size[1] y2 = y2 / config["img_h"] * origin_size[1] w = x2 - x1 h = y2 - y1 coco_results.append({ "image_id": image_id, "category_id": index2category[str(int(cls_pred.item()))], "bbox": (float(x1), float(y1), float(w), float(h)), "score": float(conf), }) logging.info("Now {}/{}".format(step, len(dataloader))) save_results_path = "coco_results.json" with open(save_results_path, "w") as f: json.dump(coco_results, f, sort_keys=True, indent=4, separators=(',', ':')) logging.info("Save coco format results to {}".format(save_results_path)) # COCO api logging.info("Using coco-evaluate tools to evaluate.") cocoGt = COCO(config["annotation_path"]) cocoDt = cocoGt.loadRes(save_results_path) cocoEval = COCOeval(cocoGt, cocoDt, "bbox") cocoEval.params.imgIds = list(coco_img_ids) # real imgIds cocoEval.evaluate() cocoEval.accumulate() cocoEval.summarize()
def test(config): is_training = False anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("load checkpoint from {}".format(config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append(YOLOLayer(config["batch_size"],i,config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # prepare images path images_path = os.listdir(config["images_path"]) images_path = [file for file in images_path if file.endswith('.jpg')] # images_path = [os.path.join(config["images_path"], name) for name in images_name] if len(images_path) == 0: raise Exception("no image found in {}".format(config["images_path"])) # Start inference batch_size = config["batch_size"] bgimage = cv2.imread(os.path.join(config["images_path"], images_path[0]), cv2.IMREAD_COLOR) bgimage = cv2.cvtColor(bgimage, cv2.COLOR_BGR2GRAY) for step in range(0, len(images_path)-1, batch_size): # preprocess images = [] images_origin = [] for path in images_path[step*batch_size: (step+1)*batch_size]: if not path.endswith(".jpg") and (not path.endswith(".png")) and not path.endswith(".JPEG"): continue image = cv2.imread(os.path.join(config["images_path"], path), cv2.IMREAD_COLOR) if image is None: logging.error("read path error: {}. skip it.".format(path)) continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) images_origin.append(image) # keep for save result image = cv2.resize(image, (config["img_w"], config["img_h"]), interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) images = np.asarray(images) images = torch.from_numpy(images).cuda() # inference with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression(output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"]) # write result images. Draw bounding boxes and labels of detections classes = open(config["classes_names_path"], "r").read().split("\n")[:-1] for idx, detections in enumerate(batch_detections): image_show =images_origin[idx] if detections is not None: anno = savexml.GEN_Annotations(path + '.jpg') anno.set_size(1280, 720, 3) unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) bbox_list = [] for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: # Rescale coordinates to original dimensions ori_h, ori_w = images_origin[idx].shape[:2] pre_h, pre_w = config["img_h"], config["img_w"] box_h = ((y2 - y1) / pre_h) * ori_h box_w = ((x2 - x1) / pre_w) * ori_w y1 = (y1 / pre_h) * ori_h x1 = (x1 / pre_w) * ori_w # Create a Rectangle patch bbox_list.append((x1, y1,box_w,box_h)) image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 0, 255), 1) boundbox = bg_judge(images_origin[idx],bgimage,bbox_list) print('boundbox',boundbox,bbox_list) for (x,y,w,h) in boundbox: image_show = cv2.rectangle(image_show, (x, y), (x + w, y + h), (0, 255, 0), 1) # anno.add_pic_attr("mouse", int(x1.cpu().data), int(y1.cpu().data), int(box_w.cpu().data) , int(box_h.cpu().data) ,"0") # # xml_path = os.path.join(config["images_path"], path).replace('rec_pic',r'detect_pic1\Annotations').replace('jpg','xml') # anno.savefile(xml_path) # cv2.imwrite(os.path.join(config["images_path"], path).replace('rec_pic',r'detect_pic1\rec_pic'),images_origin[idx]) cv2.imshow('1', image_show) cv2.waitKey(1) logging.info("Save all results to ./output/")
def evaluate(config): is_training = False # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) else: logging.warning("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append(YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader dataloader = torch.utils.data.DataLoader(dataset=COCODataset(config["test_path"], config["img_w"]), batch_size=config["batch_size"], shuffle=False, num_workers=8, pin_memory=False) # Start the eval loop #logging.info("Start eval.") n_gt = 0 correct = 0 #logging.debug('%s' % str(dataloader)) gt_histro={} pred_histro = {} correct_histro = {} for i in range(config["yolo"]["classes"]): gt_histro[i] = 1 pred_histro[i] = 1 correct_histro[i] = 0 # images 是一个batch里的全部图片,labels是一个batch里面的全部标签 for step, (images, labels) in enumerate(dataloader): labels = labels.cuda() with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) # 把三个尺度上的预测结果在第1维度(第0维度是batch里的照片,第1维度是一张照片里面的各个预测框,第2维度是各个预测数值)上拼接起来 output = torch.cat(output_list, dim=1) #logging.info('%s' % str(output.shape)) # 进行NMS抑制 #output = non_max_suppression(prediction=output, num_classes=config["yolo"]["classes"], conf_thres=config["conf_thresh"], nms_thres=config["nms_thresh"]) output = class_nms(prediction=output, num_classes=config["yolo"]["classes"],conf_thres=config["conf_thresh"], nms_thres=config["nms_thresh"]) # calculate for sample_i in range(labels.size(0)): # 计算所有的预测数量 sample_pred = output[sample_i] if sample_pred is not None: #logging.debug(sample_pred.shape) for i in range(sample_pred.shape[0]): pred_histro[int(sample_pred[i,6])] += 1 # Get labels for sample where width is not zero (dummies) target_sample = labels[sample_i, labels[sample_i, :, 3] != 0] # Ground truth的 分类编号obj_cls、相对中心x、相对中心y、相对宽w、相对高h n_gt=0 correct=0 for obj_cls, tx, ty, tw, th in target_sample: # Get rescaled gt coordinates # 转化为输入像素尺寸的 左上角像素tx1 ty1,右下角像素tx2 ty2 tx1, tx2 = config["img_w"] * (tx - tw / 2), config["img_w"] * (tx + tw / 2) ty1, ty2 = config["img_h"] * (ty - th / 2), config["img_h"] * (ty + th / 2) # 计算ground truth数量,用于统计信息 n_gt += 1 gt_histro[int(obj_cls)] += 1 # 转化为 shape(1,4)的tensor,用来计算IoU box_gt = torch.cat([coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2]]).view(1, -1) # logging.info('%s' % str(box_gt.shape)) sample_pred = output[sample_i] if sample_pred is not None: # Iterate through predictions where the class predicted is same as gt # 对于每一个ground truth,遍历预测结果 for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[sample_pred[:, 6] == obj_cls]: # 如果当前预测分类 == 当前真实分类 #logging.info("%d" % obj_cls) box_pred = torch.cat([coord.unsqueeze(0) for coord in [x1, y1, x2, y2]]).view(1, -1) #pred_histro[int(obj_pred)] += 1 iou = bbox_iou(box_pred, box_gt) #if iou >= config["iou_thres"] and obj_conf >= config["obj_thresh"]: if iou >= config["iou_thresh"]: correct += 1 correct_histro[int(obj_pred)] += 1 break #logging.debug("----------------") #logging.debug(correct_histro[4]) #logging.debug(pred_histro[4]) #logging.debug(gt_histro[4]) if n_gt: types = config["types"] reverse_types = {} # 建立一个反向的types for key in types.keys(): reverse_types[types[key]] = key #logging.info('Batch [%d/%d] mAP: %.5f' % (step, len(dataloader), float(correct / n_gt))) logging.info('Precision:%s' % str([reverse_types[i] +':'+ str(int(100 * correct_histro[i] / pred_histro[i])) for i in range(config["yolo"]["classes"]) ])) logging.info('Recall :%s' % str([reverse_types[i] +':'+ str(int(100 * correct_histro[i] / gt_histro[i])) for i in range(config["yolo"]["classes"])]))
def test(config): is_training = False anchors = [int(x) for x in config["yolo"]["anchors"].split(",")] anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i in range(0, len(anchors), 6)] anchors.reverse() config["yolo"]["anchors"] = [] for i in range(3): config["yolo"]["anchors"].append(anchors[i]) net = ModelMain(config, is_training=is_training) net.train(is_training) # Set data parallel net = nn.DataParallel(net) net = net.cuda() ini_files = [ inifile for inifile in os.listdir( os.path.join(config['test_weights'], 'result')) if inifile.endswith('.ini') ] accuracy_s = [(inifile[:-4]).split('_')[-1] for inifile in ini_files] accuracy_ints = list(map(float, accuracy_s)) max_index = accuracy_ints.index(max(accuracy_ints)) # for kkk,ini_file in enumerate(ini_files): ini_list_config = configparser.ConfigParser() config_file_path = os.path.join(config['test_weights'], 'result', ini_files[max_index]) Bi_picpath = os.path.join(config['test_weights'], 'result', ini_files[max_index]).replace('.ini', '') os.makedirs(Bi_picpath, exist_ok=True) ini_list_config.read(config_file_path) ini_session = ini_list_config.sections() accuracy = ini_list_config.items(ini_session[0]) err_jpgfiles = ini_list_config.items(ini_session[1]) weight_file = os.path.join( config['test_weights'], '%s.weights' % ini_files[max_index].split('_')[0]) if weight_file: # Restore pretrain model logging.info("load checkpoint from {}".format(weight_file)) state_dict = torch.load(weight_file) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # images_name = os.listdir(config["images_path"]) # prepare images path # images_path = [os.path.join(config["images_path"], name) for name in images_name] # if len(images_path) == 0: # raise Exception("no image found in {}".format(config["images_path"])) # batch_size = config["batch_size"]# Start inference # for step in range(0, len(images_path), batch_size): for _jpg_images in err_jpgfiles: images = [] # preprocess images_origin = [] jpg_path = str(_jpg_images[1]) logging.info("processing: {}".format(jpg_path)) bbox_list = read_gt_boxes(jpg_path) image = cv2.imread(jpg_path, cv2.IMREAD_COLOR) if image is None: logging.error("read path error: {}. skip it.".format(jpg_path)) continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) images_origin.append(image) # keep for save result image = cv2.resize(image, (config["img_w"], config["img_h"]), interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) images = np.asarray(images) images = torch.from_numpy(images).cuda() with torch.no_grad(): # inference outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression( output, config["yolo"]["classes"], conf_thres=config["confidence_threshold"]) classes = open(config["classes_names_path"], "r").read().split("\n")[:-1] for idx, detections in enumerate(batch_detections): image_show = images_origin[idx] if detections is not None: unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: ori_h, ori_w = images_origin[ idx].shape[: 2] # Rescale coordinates to original dimensions pre_h, pre_w = config["img_h"], config["img_w"] box_h = ((y2 - y1) / pre_h) * ori_h box_w = ((x2 - x1) / pre_w) * ori_w y1 = (y1 / pre_h) * ori_h x1 = (x1 / pre_w) * ori_w image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 255, 0), 2) for (x1, x2, y1, y2) in bbox_list: [x1, x2, y1, y2] = map(int, [x1, x2, y1, y2]) cv2.rectangle(image_show, (x1, y1), (x2, y2), (255, 0, 0), 2) pic_name = (jpg_path.split('/')[-1]).split('.')[0] image_show = cv2.cvtColor(image_show, cv2.COLOR_RGB2BGR) cv2.imwrite( os.path.join(Bi_picpath, '%s.jpg' % os.path.basename(pic_name)), image_show)
def test(): is_traning = False # 不训练,测试 # Load and initialize network net = ModelMain(is_training=is_traning) net.train(is_traning) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if trained_model_dir: logging.info("load checkpoint from {}".format(trained_model_dir)) state_dict = torch.load(trained_model_dir) net.load_state_dict(state_dict) else: raise Exception("missing pretrain_snapshot!!!") # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append(YOLOLoss(anchors[i], yolo_class_num, (img_h, img_w))) # prepare the images path images_name = os.listdir("./images/") images_path = [os.path.join("./images/", name) for name in images_name] print('images_name:', images_name) print('images_path:', len(images_path), images_path) if len(images_path) == 0: raise Exception("no image found in {}".format("./images/")) # Start inference batch_size = 16 for step in range(0, len(images_path), batch_size): # range(0, 4, 16) step = 0, 4, 8, 12 logging.info('Batch_size:{}'.format(batch_size)) # preprocess images = [] # 输入网络图片组 images_origin = [] # 原始图片组 for path in images_path[step * batch_size:(step + 1) * batch_size]: logging.info("processing: {}".format(path)) image = cv2.imread(path, cv2.IMREAD_COLOR) # cv2.imshow('Image', image) # cv2.waitKey(0) logging.info(" √ Successfully Processed! √") if image is None: logging.error("read path error: {}. skip it.".format(path)) continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) images_origin.append(image) # 预处理完毕,加入原始图片组 # 进一步将图片处理为网络可以接受的数据类型(resize、归一化等) image = cv2.resize(image, (img_h, img_w), interpolation=cv2.INTER_LINEAR) image = image.astype(np.float32) image /= 255.0 image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) images.append(image) # 归一化完毕,加入输入网络图片组 images = np.asarray(images) images = torch.from_numpy(images).cuda() logging.info("\nImages Convert to Tensor of CUDA Done!") # inference with torch.no_grad(): outputs = net(images) output_list = [] for i in range(3): output_list.append(yolo_losses[i](outputs[i])) output = torch.cat(output_list, 1) batch_detections = non_max_suppression(prediction=output, num_classes=yolo_class_num, conf_thres=0.5) logging.info("\nNet Detection Done!\n") # write result images: Draw BBox classes = open(classes_name_path, 'r').read().split("\n")[:-1] # 读取coco.names if not os.path.isdir("./output/"): os.makedirs("./output/") for idx, detections in enumerate(batch_detections): plt.figure() fig, ax = plt.subplots(1) ax.imshow(images_origin[idx]) if detections is not None: unique_labels = detections[:, -1].cpu().unique() n_cls_preds = len(unique_labels) bbox_colors = random.sample(colors, n_cls_preds) # print('Final Detections: ', detections) for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: color = bbox_colors[int( np.where(unique_labels == int(cls_pred))[0])] # Rescale coordinates to original dimensions ori_h, ori_w = images_origin[idx].shape[:2] pre_h, pre_w = img_h, img_w # 416, 416 box_h = ((y2 - y1) / pre_h) * ori_h box_w = ((x2 - x1) / pre_w) * ori_w y1 = (y1 / pre_h) * ori_h x1 = (x1 / pre_w) * ori_w # Create a Rectangle patch bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolor='none') # Add the bbox to the plot ax.add_patch(bbox) # Add label plt.text(x1, y1, s=classes[int(cls_pred)], color='white', verticalalignment='top', bbox={ 'color': color, 'pad': 0 }) # Save generated image with detections plt.axis('off') plt.gca().xaxis.set_major_locator(NullLocator()) plt.gca().yaxis.set_major_locator(NullLocator()) plt.savefig('output/{}_{}.jpg'.format(step, idx), bbox_inches='tight', pad_inches=0.0) plt.close() logging.info("All the Test Process Succeed! Enjoy it!")
def train(config): config["global_step"] = config.get("start_step", 0) is_training = False if config.get("export_onnx") else True # Load and initialize network net = ModelMain(config, is_training=is_training) net.train(is_training) # Optimizer and learning rate optimizer = _get_optimizer(config, net) lr_scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=config["lr"]["decay_step"], gamma=config["lr"]["decay_gamma"]) # Set data parallel net = nn.DataParallel(net) net = net.cuda() # Restore pretrain model if config["pretrain_snapshot"]: logging.info("Load pretrained weights from {}".format( config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) # Only export onnx # if config.get("export_onnx"): # real_model = net.module # real_model.eval() # dummy_input = torch.randn(8, 3, config["img_h"], config["img_w"]).cuda() # save_path = os.path.join(config["sub_working_dir"], "pytorch.onnx") # logging.info("Exporting onnx to {}".format(save_path)) # torch.onnx.export(real_model, dummy_input, save_path, verbose=False) # logging.info("Done. Exiting now.") # sys.exit() # Evaluate interface # if config["evaluate_type"]: # logging.info("Using {} to evaluate model.".format(config["evaluate_type"])) # evaluate_func = importlib.import_module(config["evaluate_type"]).run_eval # config["online_net"] = net # YOLO loss with 3 scales yolo_losses = [] for i in range(3): yolo_losses.append( YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"]))) # DataLoader dataloader = torch.utils.data.DataLoader(COCODataset( config["train_path"], (config["img_w"], config["img_h"]), is_training=True), batch_size=config["batch_size"], shuffle=True, num_workers=32, pin_memory=True) # Start the training loop logging.info("Start training.") for epoch in range(config["epochs"]): for step, samples in enumerate(dataloader): images, labels = samples["image"], samples["label"] start_time = time.time() config["global_step"] += 1 # Forward and backward optimizer.zero_grad() outputs = net(images) losses_name = ["total_loss", "x", "y", "w", "h", "conf", "cls"] losses = [] for _ in range(len(losses_name)): losses.append([]) for i in range(3): _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): losses[j].append(l) losses = [sum(l) for l in losses] loss = losses[0] loss.backward() optimizer.step() if step > 0 and step % 10 == 0: _loss = loss.item() duration = float(time.time() - start_time) example_per_second = config["batch_size"] / duration lr = optimizer.param_groups[0]['lr'] logging.info( "epoch [%.3d] iter = %d loss = %.2f example/sec = %.3f lr = %.5f " % (epoch, step, _loss, example_per_second, lr)) config["tensorboard_writer"].add_scalar( "lr", lr, config["global_step"]) config["tensorboard_writer"].add_scalar( "example/sec", example_per_second, config["global_step"]) for i, name in enumerate(losses_name): value = _loss if i == 0 else losses[i] config["tensorboard_writer"].add_scalar( name, value, config["global_step"]) # if step > 0 and step % 1000 == 0: # net.train(False) # _save_checkpoint(net.state_dict(), config) # net.train(True) _save_checkpoint(net.state_dict(), config) lr_scheduler.step() # net.train(False) _save_checkpoint(net.state_dict(), config) # net.train(True) logging.info("Bye~")
def train(config): # Hyper-parameters config["global_step"] = config.get("start_step", 0) is_training = True # Net & Loss & Optimizer ## Net Main net = ModelMain(config, is_training=is_training) net.train(is_training) ## YOLO Loss with 3 scales yolo_losses = [] for i in range(3): yolo_loss = YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"], (config["img_w"], config["img_h"])) yolo_losses.append(yolo_loss) ## Optimizer and LR scheduler optimizer = _get_optimizer(config, net) lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=config["lr"]["decay_step"], gamma=config["lr"]["decay_gamma"]) net = nn.DataParallel(net) net = net.cuda() # Load checkpoint if config["pretrain_snapshot"]: logging.info("Load pretrained weights from {}".format(config["pretrain_snapshot"])) state_dict = torch.load(config["pretrain_snapshot"]) net.load_state_dict(state_dict) # DataLoader dataloader = torch.utils.data.DataLoader(AIPrimeDataset(config["train_path"]), batch_size=config["batch_size"], shuffle=True, num_workers=16, pin_memory=False) # Start the training logging.info("Start training.") for epoch in range(config["start_epoch"], config["epochs"]): for step, (images, labels) in enumerate(dataloader): start_time = time.time() config["global_step"] += 1 # Forward outputs = net(images) # Loss losses_name = ["total_loss", "x", "y", "w", "h", "conf", "cls"] losses = [[]] * len(losses_name) for i in range(3): _loss_item = yolo_losses[i](outputs[i], labels) for j, l in enumerate(_loss_item): losses[j].append(l) losses = [sum(l) for l in losses] loss = losses[0] # Zero & Backward & Step optimizer.zero_grad() loss.backward() optimizer.step() # Logging if step > 0 and step % 10 == 0: _loss = loss.item() duration = float(time.time() - start_time) example_per_second = config["batch_size"] / duration lr = optimizer.param_groups[0]['lr'] logging.info( "epoch [%.3d] iter = %d loss = %.2f example/sec = %.3f lr = %.5f " % (epoch, step, _loss, example_per_second, lr) ) # Things to be done for every epoch ## LR schedule lr_scheduler.step() ## Save checkpoint _save_checkpoint(net.state_dict(), config, epoch) # Finish training logging.info("QiaJiaBa~ BeiBei")