def main(): """ train dataloader """ args = parser.parse_args() data_loader = TrainDataLoader(args.train_path, check=True) if not os.path.exists(args.weight_dir): os.makedirs(args.weight_dir) """ compute max_batches """ for root, dirs, files in os.walk(args.train_path): for dirname in dirs: dir_path = os.path.join(root, dirname) args.max_batches += len(os.listdir(dir_path)) inter = args.max_batches // 10 print('Max batches:{} in one epoch '.format(args.max_batches)) """ Model on gpu """ model = SiameseRPN() model = model.cuda() cudnn.benchmark = True """ loss and optimizer """ criterion = MultiBoxLoss() optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) """ load weights """ init_weights(model) if not args.checkpoint_path == None: assert os.path.isfile( args.checkpoint_path), '{} is not valid checkpoint_path'.format( args.checkpoint_path) try: checkpoint = torch.load(args.checkpoint_path) start = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) except: start = 0 init_weights(model) else: start = 0 """ train phase """ closses, rlosses, tlosses = AverageMeter(), AverageMeter(), AverageMeter() for epoch in range(start, args.max_epoches): cur_lr = adjust_learning_rate(args.lr, optimizer, epoch, gamma=0.1) index_list = range(data_loader.__len__()) #for example in range(args.max_batches): for example in range(900): ret = data_loader.__get__(random.choice(index_list)) template = ret['template_tensor'].cuda() detection = ret['detection_tensor'].cuda() pos_neg_diff = ret['pos_neg_diff_tensor'].cuda( ) if ret['pos_neg_diff_tensor'] is not None else None cout, rout = model(template, detection) predictions = (cout, rout) targets = pos_neg_diff area = ret['area_target_in_resized_detection'] num_pos = len(np.where(pos_neg_diff == 1)[0]) if area == 0 or num_pos == 0 or pos_neg_diff is None: continue closs, rloss, loss, reg_pred, reg_target, pos_index, neg_index = criterion( predictions, targets) # debug for class cout = cout.squeeze().permute(1, 2, 0).reshape(-1, 2) cout = cout.cpu().detach().numpy() print(cout.shape) score = 1 / (1 + np.exp(cout[:, 0] - cout[:, 1])) print(score[pos_index]) print(score[neg_index]) #time.sleep(1) # debug for reg tmp_dir = '/home/song/srpn/tmp/visualization/7_train_debug_pos_anchors' if not os.path.exists(tmp_dir): os.makedirs(tmp_dir) detection = ret['detection_cropped_resized'].copy() draw = ImageDraw.Draw(detection) pos_anchors = ret['pos_anchors'].copy() # pos anchor的回归情况 x = pos_anchors[:, 0] + pos_anchors[:, 2] * reg_pred[ pos_index, 0].cpu().detach().numpy() y = pos_anchors[:, 1] + pos_anchors[:, 3] * reg_pred[ pos_index, 1].cpu().detach().numpy() w = pos_anchors[:, 2] * np.exp(reg_pred[pos_index, 2].cpu().detach().numpy()) h = pos_anchors[:, 3] * np.exp(reg_pred[pos_index, 3].cpu().detach().numpy()) x1s, y1s, x2s, y2s = x - w // 2, y - h // 2, x + w // 2, y + h // 2 for i in range(2): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='red') #predict # 应当的gt x = pos_anchors[:, 0] + pos_anchors[:, 2] * reg_target[ pos_index, 0].cpu().detach().numpy() y = pos_anchors[:, 1] + pos_anchors[:, 3] * reg_target[ pos_index, 1].cpu().detach().numpy() w = pos_anchors[:, 2] * np.exp( reg_target[pos_index, 2].cpu().detach().numpy()) h = pos_anchors[:, 3] * np.exp( reg_target[pos_index, 3].cpu().detach().numpy()) x1s, y1s, x2s, y2s = x - w // 2, y - h // 2, x + w // 2, y + h // 2 for i in range(2): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='green') #gt # 找分数zui da de, m_indexs = np.argsort(score)[::-1][:5] for m_index in m_indexs: diff = reg_pred[m_index].cpu().detach().numpy() anc = ret['anchors'][m_index] x = anc[0] + anc[0] * diff[0] y = anc[1] + anc[1] * diff[1] w = anc[2] * np.exp(diff[2]) h = anc[3] * np.exp(diff[3]) x1, y1, x2, y2 = x - w // 2, y - h // 2, x + w // 2, y + h // 2 draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=2, fill='black') save_path = osp.join( tmp_dir, 'epoch_{:04d}_{:04d}_{:02d}.jpg'.format(epoch, example, i)) detection.save(save_path) closs_ = closs.cpu().item() if np.isnan(closs_): sys.exit(0) #loss = closs + rloss closses.update(closs.cpu().item()) rlosses.update(rloss.cpu().item()) tlosses.update(loss.cpu().item()) optimizer.zero_grad() loss.backward() optimizer.step() #time.sleep(1) print( "Epoch:{:04d}\texample:{:08d}/{:08d}({:.2f})\tlr:{:.7f}\tcloss:{:.4f}\trloss:{:.4f}\ttloss:{:.4f}" .format(epoch, example + 1, args.max_batches, 100 * (example + 1) / args.max_batches, cur_lr, closses.avg, rlosses.avg, tlosses.avg)) if epoch % 5 == 0: file_path = os.path.join( args.weight_dir, 'epoch_{:04d}_weights.pth.tar'.format(epoch)) state = { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), } torch.save(state, file_path)
def main(): """ train dataloader """ args = parser.parse_args() data_loader = TrainDataLoader(args.train_path, check = args.debug) if not os.path.exists(args.weight_dir): os.makedirs(args.weight_dir) """ compute max_batches """ for root, dirs, files in os.walk(args.train_path): for dirname in dirs: dir_path = os.path.join(root, dirname) args.max_batches += len(os.listdir(dir_path)) print('max_batches: {}'.format(args.max_batches)) """ Model on gpu """ model = SiameseRPN() model = model.cuda() cudnn.benchmark = True """ loss and optimizer """ criterion = MultiBoxLoss() optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay = args.weight_decay) """ load weights """ init_weights(model) if not args.checkpoint_path == None: assert os.path.isfile(args.checkpoint_path), '{} is not valid checkpoint_path'.format(args.checkpoint_path) try: checkpoint = torch.load(args.checkpoint_path) start = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) except: start = 0 init_weights(model) else: start = 0 """ train phase """ closses, rlosses, tlosses = AverageMeter(), AverageMeter(), AverageMeter() steps = 0 #print('data_loader length: {}'.format(len(data_loader))) for epoch in range(start, args.max_epoches): cur_lr = adjust_learning_rate(args.lr, optimizer, epoch, gamma=0.1) index_list = range(data_loader.__len__()) example_index = 0 for example in range(args.max_batches): ret = data_loader.__get__(random.choice(index_list)) template = ret['template_tensor'].cuda() detection= ret['detection_tensor'].cuda() pos_neg_diff = ret['pos_neg_diff_tensor'].cuda() cout, rout = model(template, detection) predictions, targets = (cout, rout), pos_neg_diff closs, rloss, loss, reg_pred, reg_target, pos_index, neg_index = criterion(predictions, targets) closs_ = closs.cpu().item() if np.isnan(closs_): sys.exit(0) closses.update(closs.cpu().item()) rlosses.update(rloss.cpu().item()) tlosses.update(loss.cpu().item()) optimizer.zero_grad() loss.backward() optimizer.step() steps += 1 cout = cout.cpu().detach().numpy() score = 1/(1 + np.exp(cout[:,0]-cout[:,1])) # ++++++++++++ post process below just for debug ++++++++++++++++++++++++ # ++++++++++++++++++++ v1.0 add penalty +++++++++++++++++++++++++++++++++ if ret['pos_anchors'] is not None: penalty_k = 0.055 tx, ty, tw, th = ret['template_target_xywh'].copy() tw *= ret['template_cropprd_resized_ratio'] th *= ret['template_cropprd_resized_ratio'] anchors = ret['anchors'].copy() w = anchors[:,2] * np.exp(reg_pred[:, 2].cpu().detach().numpy()) h = anchors[:,3] * np.exp(reg_pred[:, 3].cpu().detach().numpy()) eps = 1e-2 change_w = np.maximum(w/(tw+eps), tw/(w+eps)) change_h = np.maximum(h/(th+eps), th/(h+eps)) penalty = np.exp(-(change_w + change_h - 1) * penalty_k) pscore = score * penalty else: pscore = score # +++++++++++++++++++ v1.0 add window default cosine ++++++++++++++++++++++ score_size = 17 window_influence = 0.42 window = (np.outer(np.hanning(score_size), np.hanning(score_size)).reshape(17,17,1) + np.zeros((1, 1, 5))).reshape(-1) pscore = pscore * (1 - window_influence) + window * window_influence score_old = score score = pscore #from 0.2 - 0.7 # ++++++++++++++++++++ debug for class ++++++++++++++++++++++++++++++++++++ if example_index%1000 == 0: print(score[pos_index]) # this should tend to be 1 print(score[neg_index]) # this should tend to be 0 # ++++++++++++++++++++ debug for reg ++++++++++++++++++++++++++++++++++++++ tmp_dir = '/home/ly/chz/Siamese-RPN-pytorch/code_v1.0/tmp/visualization/7_check_train_phase_debug_pos_anchors' if not os.path.exists(tmp_dir): os.makedirs(tmp_dir) detection = ret['detection_cropped_resized'].copy() draw = ImageDraw.Draw(detection) pos_anchors = ret['pos_anchors'].copy() if ret['pos_anchors'] is not None else None if pos_anchors is not None: # draw pos anchors x = pos_anchors[:, 0] y = pos_anchors[:, 1] w = pos_anchors[:, 2] h = pos_anchors[:, 3] x1s, y1s, x2s, y2s = x - w//2, y - h//2, x + w//2, y + h//2 for i in range(16): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='white') # pos anchor # pos anchor transform to red box after prediction x = pos_anchors[:,0] + pos_anchors[:, 2] * reg_pred[pos_index, 0].cpu().detach().numpy() y = pos_anchors[:,1] + pos_anchors[:, 3] * reg_pred[pos_index, 1].cpu().detach().numpy() w = pos_anchors[:,2] * np.exp(reg_pred[pos_index, 2].cpu().detach().numpy()) h = pos_anchors[:,3] * np.exp(reg_pred[pos_index, 3].cpu().detach().numpy()) x1s, y1s, x2s, y2s = x - w//2, y - h//2, x + w//2, y + h//2 for i in range(16): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='red') # predict(white -> red) # pos anchor should be transformed to green gt box, if red and green is same, it is overfitting x = pos_anchors[:,0] + pos_anchors[:, 2] * reg_target[pos_index, 0].cpu().detach().numpy() y = pos_anchors[:,1] + pos_anchors[:, 3] * reg_target[pos_index, 1].cpu().detach().numpy() w = pos_anchors[:,2] * np.exp(reg_target[pos_index, 2].cpu().detach().numpy()) h = pos_anchors[:,3] * np.exp(reg_target[pos_index, 3].cpu().detach().numpy()) x1s, y1s, x2s, y2s = x - w//2, y-h//2, x + w//2, y + h//2 for i in range(16): x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i] draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='green') # gt (white -> green) x1, y1, x3, y3 = x1s[0], y1s[0], x2s[0], y2s[0] else: x1, y1, x3, y3 = 0, 0, 0, 0 # top1 proposal after nms (white) if example_index%1000 == 0: save_path = osp.join(tmp_dir, 'epoch_{:010d}_{:010d}_anchor_pred.jpg'.format(epoch, example)) detection.save(save_path) example_index = example_index+1 # +++++++++++++++++++ v1.0 restore ++++++++++++++++++++++++++++++++++++++++ ratio = ret['detection_cropped_resized_ratio'] detection_cropped = ret['detection_cropped'].copy() detection_cropped_resized = ret['detection_cropped_resized'].copy() original = Image.open(ret['detection_img_path']) x_, y_ = ret['detection_tlcords_of_original_image'] draw = ImageDraw.Draw(original) w, h = original.size """ un resized """ x1, y1, x3, y3 = x1/ratio, y1/ratio, y3/ratio, y3/ratio """ un cropped """ x1 = np.clip(x_ + x1, 0, w-1).astype(np.int32) # uncropped #target_of_original_img y1 = np.clip(y_ + y1, 0, h-1).astype(np.int32) x3 = np.clip(x_ + x3, 0, w-1).astype(np.int32) y3 = np.clip(y_ + y3, 0, h-1).astype(np.int32) draw.line([(x1, y1), (x3, y1), (x3, y3), (x1, y3), (x1, y1)], width=3, fill='yellow') #save_path = osp.join(tmp_dir, 'epoch_{:010d}_{:010d}_restore.jpg'.format(epoch, example)) #original.save(save_path) print("Epoch:{:04d}\texample:{:06d}/{:06d}({:.2f})%\tsteps:{:010d}\tlr:{:.7f}\tcloss:{:.4f}\trloss:{:.4f}\ttloss:{:.4f}".format(epoch, example+1, args.max_batches, 100*(example+1)/args.max_batches, steps, cur_lr, closses.avg, rlosses.avg, tlosses.avg )) if steps % 1 == 0: file_path = os.path.join(args.weight_dir, 'weights-{:07d}.pth.tar'.format(steps)) state = { 'epoch' :epoch+1, 'state_dict' :model.state_dict(), 'optimizer' : optimizer.state_dict(), } torch.save(state, file_path)
def main(): """ dataloader """ args = parser.parse_args() data_loader = TrainDataLoader(args.test_path, out_feature=25, check=False) """ Model on gpu """ model = SiameseRPN() model = model.cuda() cudnn.benchmark = True """ loss and optimizer """ criterion = MultiBoxLoss() """ load weights """ init_weights(model) if args.checkpoint_path == None: sys.exit('please input trained model') else: assert os.path.isfile( args.checkpoint_path), '{} is not valid checkpoint_path'.format( args.checkpoint_path) checkpoint = torch.load(args.checkpoint_path) start = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) """ test phase """ index_list = range(data_loader.__len__()) threshold = 50 precision = [] precision_c = [] average_error = [] average_error_c = [] iou = [] iou_c = [] for example in range(args.max_batches): ret = data_loader.__get__(random.choice(index_list)) template = ret['template_tensor'].cuda() detection = ret['detection_tensor'].cuda() pos_neg_diff = ret['pos_neg_diff_tensor'].cuda() cout, rout = model(template, detection) #[1, 10, 17, 17], [1, 20, 17, 17] template_img = ret['template_cropped_transformed'] detection_img = ret['detection_cropped_transformed'] cout = cout.reshape(-1, 2) rout = rout.reshape(-1, 4) cout = cout.cpu().detach().numpy() score = 1 / (1 + np.exp(cout[:, 0] - cout[:, 1])) diff = rout.cpu().detach().numpy() #1445 num_proposals = 1 score_64_index = np.argsort(score)[::-1][:num_proposals] score64 = score[score_64_index] diffs64 = diff[score_64_index, :] anchors64 = ret['anchors'][score_64_index] proposals_x = (anchors64[:, 0] + anchors64[:, 2] * diffs64[:, 0]).reshape(-1, 1) proposals_y = (anchors64[:, 1] + anchors64[:, 3] * diffs64[:, 1]).reshape(-1, 1) proposals_w = (anchors64[:, 2] * np.exp(diffs64[:, 2])).reshape(-1, 1) proposals_h = (anchors64[:, 3] * np.exp(diffs64[:, 3])).reshape(-1, 1) proposals = np.hstack( (proposals_x, proposals_y, proposals_w, proposals_h)) d = os.path.join(ret['tmp_dir'], '6_pred_proposals') if not os.path.exists(d): os.makedirs(d) template = ret['template_cropped_transformed'] save_path = os.path.join(ret['tmp_dir'], '6_pred_proposals', '{:04d}_0_template.jpg'.format(example)) template.save(save_path) """traditional correlation match method""" template_img = cv2.cvtColor(np.asarray(template_img), cv2.COLOR_RGB2BGR) detection_img = cv2.cvtColor(np.asarray(detection_img), cv2.COLOR_RGB2BGR) res = cv2.matchTemplate(detection_img, template_img, cv2.TM_CCOEFF_NORMED) min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res) x1_c = max_loc[0] y1_c = max_loc[1] """ visualization """ ratio = ret['detection_cropped_resized_ratio'] original = Image.open(ret['detection_img_path']) origin_w, origin_h = original.size x_, y_ = ret['detection_tlcords_of_original_image'] draw = ImageDraw.Draw(original) for i in range(num_proposals): x, y, w, h = proposals_x[i], proposals_y[i], proposals_w[ i], proposals_h[i] x1, y1, x3, y3 = x - w // 2, y - h // 2, x + w // 2, y + h // 2 """ un resized """ x1, y1, x3, y3 = x1 / ratio, y1 / ratio, x3 / ratio, y3 / ratio x1_c, y1_c = x1_c / ratio, y1_c / ratio """ un cropped """ x1_g, y1_g, w, h = ret['template_target_x1y1wh'] x3_g = x1_g + w y3_g = y1_g + h x1 = np.clip(x_ + x1, 0, origin_w - 1).astype( np.int32) # uncropped #target_of_original_img y1 = np.clip(y_ + y1, 0, origin_h - 1).astype(np.int32) x3 = np.clip(x_ + x3, 0, origin_w - 1).astype(np.int32) y3 = np.clip(y_ + y3, 0, origin_h - 1).astype(np.int32) x1_c = np.clip(x_ + x1_c, 0, origin_w - 1).astype(np.int32) y1_c = np.clip(y_ + y1_c, 0, origin_h - 1).astype(np.int32) x3_c = x1_c + ret['template_target_xywh'][2] y3_c = y1_c + ret['template_target_xywh'][3] draw.line([(x1, y1), (x3, y1), (x3, y3), (x1, y3), (x1, y1)], width=3, fill='yellow') draw.line([(x1_g, y1_g), (x3_g, y1_g), (x3_g, y3_g), (x1_g, y3_g), (x1_g, y1_g)], width=3, fill='blue') draw.line([(x1_c, y1_c), (x3_c, y1_c), (x3_c, y3_c), (x1_c, y3_c), (x1_c, y1_c)], width=3, fill='red') save_path = os.path.join(ret['tmp_dir'], '6_pred_proposals', '{:04d}_1_restore.jpg'.format(example)) original.save(save_path) print('save at {}'.format(save_path)) """compute iou""" s1 = np.array([x1, y1, x3, y1, x3, y3, x1, y3, x1, y1]) s2 = np.array( [x1_g, y1_g, x3_g, y1_g, x3_g, y3_g, x1_g, y3_g, x1_g, y1_g]) s3 = np.array( [x1_c, y1_c, x3_c, y1_c, x3_c, y3_c, x1_c, y3_c, x1_c, y1_c]) iou.append(intersection(s1, s2)) iou_c.append(intersection(s3, s2)) """compute average error""" cx = (x1 + x3) / 2 cy = (y1 + y3) / 2 cx_g = (x1_g + x3_g) / 2 cy_g = (y1_g + y3_g) / 2 cx_c = (x1_c + x3_c) / 2 cy_c = (y1_c + y3_c) / 2 error = math.sqrt(math.pow(cx - cx_g, 2) + math.pow(cy - cy_g, 2)) error_c = math.sqrt(math.pow(cx - cx_c, 2) + math.pow(cy - cy_c, 2)) average_error.append(error) average_error_c.append(error_c) if error <= threshold: precision.append(1) else: precision.append(0) if error_c <= threshold: precision_c.append(1) else: precision_c.append(0) iou_mean = np.mean(np.array(iou)) error_mean = np.mean(np.array(average_error)) iou_mean_c = np.mean(np.array(iou_c)) error_mean_c = np.mean(np.array(average_error_c)) precision = np.mean(np.array(precision)) precision_c = np.mean(np.array(precision_c)) print('average iou: {:.4f}'.format(iou_mean)) print('average error: {:.4f}'.format(error_mean)) print('average iou for traditional method: {:.4f}'.format(iou_mean_c)) print('average error for traditional method: {:.4f}'.format(error_mean_c)) print('precision: {:.4f} @ threshold {:02d}'.format(precision, threshold)) print('precision for traditional method: {:.4f} @ threshold {:02d}'.format( precision_c, threshold))
def main(): """ dataloader """ args = parser.parse_args() data_loader = TrainDataLoader(args.train_path, check=False) """ Model on gpu """ model = SiameseRPN() model = model.cuda() cudnn.benchmark = True """ loss and optimizer """ criterion = MultiBoxLoss() """ load weights """ init_weights(model) if args.checkpoint_path == None: sys.exit('please input trained model') else: assert os.path.isfile( args.checkpoint_path), '{} is not valid checkpoint_path'.format( args.checkpoint_path) checkpoint = torch.load(args.checkpoint_path) start = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) """ test phase """ index_list = range(data_loader.__len__()) for example in range(args.max_batches): ret = data_loader.__get__(random.choice(index_list)) template = ret['template_tensor'].cuda() detection = ret['detection_tensor'].cuda() pos_neg_diff = ret['pos_neg_diff_tensor'].cuda() cout, rout = model(template, detection) #[1, 10, 17, 17], [1, 20, 17, 17] cout = cout.reshape(-1, 2) rout = rout.reshape(-1, 4) cout = cout.cpu().detach().numpy() print('cout size: {}'.format(cout.shape)) score = 1 / (1 + np.exp(cout[:, 1] - cout[:, 0])) print('score: {}, size: {}'.format(score, score.shape)) diff = rout.cpu().detach().numpy() #1445 num_proposals = 15 score_64_index = np.argsort(score)[::-1][:num_proposals] print('score_64_index: {}, size: {}'.format(score_64_index, score_64_index.shape)) score64 = score[score_64_index] print('score: {}'.format(score64)) diffs64 = diff[score_64_index, :] anchors64 = ret['anchors'][score_64_index] proposals_x = (anchors64[:, 0] + anchors64[:, 2] * diffs64[:, 0]).reshape(-1, 1) proposals_y = (anchors64[:, 1] + anchors64[:, 3] * diffs64[:, 1]).reshape(-1, 1) proposals_w = (anchors64[:, 2] * np.exp(diffs64[:, 2])).reshape(-1, 1) proposals_h = (anchors64[:, 3] * np.exp(diffs64[:, 3])).reshape(-1, 1) proposals = np.hstack( (proposals_x, proposals_y, proposals_w, proposals_h)) d = os.path.join(ret['tmp_dir'], '6_pred_proposals') if not os.path.exists(d): os.makedirs(d) detection = ret['detection_cropped_resized'] save_path = os.path.join(ret['tmp_dir'], '6_pred_proposals', '{:04d}_1_detection.jpg'.format(example)) detection.save(save_path) template = ret['template_cropped_resized'] save_path = os.path.join(ret['tmp_dir'], '6_pred_proposals', '{:04d}_0_template.jpg'.format(example)) template.save(save_path) """ 可视化 """ draw = ImageDraw.Draw(detection) for i in range(num_proposals): x, y, w, h = proposals_x[i], proposals_y[i], proposals_w[ i], proposals_h[i] x1, y1, x2, y2 = x - w // 2, y - h // 2, x + w // 2, y + h // 2 draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='red') """ save detection template proposals""" save_path = os.path.join(ret['tmp_dir'], '6_pred_proposals', '{:04d}_2_proposals.jpg'.format(example)) detection.save(save_path) print('save at {}'.format(save_path)) # restore """
import cv2 # imread import torch import torch.nn as nn import numpy as np import scipy.io as scio import os import time from os.path import realpath, dirname, join from net import SiameseRPN from run_SiamRPN import SiamRPN_init, SiamRPN_track from utils import get_axis_aligned_bbox, cxy_wh_2_rect if __name__ == '__main__': # load net checkpoint_path = '/home/ly/chz/weights-0690000.pth.tar' net = SiameseRPN() net = net.eval().cuda() checkpoint = torch.load(checkpoint_path) net.load_state_dict(checkpoint['state_dict']) #net_file = join(realpath(dirname(__file__)), 'SiamRPNBIG.model') #net = SiamRPNBIG() #net.load_state_dict(torch.load(net_file)) #net.eval().cuda() OTB100_path = '/home/ly/chz/OTB' result_path = '/home/ly/chz/srpn_result' if not os.path.isdir(result_path): os.mkdir(result_path)