Beispiel #1
0
def main():
    args = parser.parse_args()

    """ train dataloader """
    data_loader = TrainDataLoader("C:\\Users\\sport\\Desktop\\SiamMask-Pytorch\\DAVIS-4\\JPEGImages\\480p")

    print('-')

    """ 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))

    # Setup Model
    cfg = load_config(args)
    from experiments.siammask.custom import Custom
    model = Custom(anchors=cfg['anchors'])
    if args.resume:
        assert isfile(args.resume), '{} is not a valid file'.format(args.resume)
        model = load_pretrain(model, args.resume)

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    model.eval().to(device)
    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)

    """ train phase """
    closses, rlosses, tlosses = AverageMeter(), AverageMeter(), AverageMeter()
    steps = 0
    start = 0
    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):
            ret = data_loader.__get__(random.choice(index_list)) 
            template = ret['template_tensor'].to(device)
            detection= ret['detection_tensor'].to(device)
            mask_target = ret['mask_template_tensor'].to(device)
            pos_neg_diff = ret['pos_neg_diff_tensor'].to(device)
            cout, rout, mask = model(template, detection)
            predictions, targets = (cout, rout, mask), pos_neg_diff
            closs, rloss, mloss, loss, reg_pred, reg_target, pos_index, neg_index = criterion(predictions, targets, mask_target)
            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]))
            print("Epoch:{:04d}\texample:{:06d}/{:06d}({:.2f})%\tsteps:{:010d}\tlr:{:.7f}\tcloss:ss:{:.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 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)
Beispiel #2
0
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)
Beispiel #3
0
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)
Beispiel #4
0
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
        """
def main():
    """parameter initialization"""
    args = parser.parse_args()
    exp_name_dir = experiment_name_dir(args.experiment_name)
    """Load the parameters from json file"""
    json_path = os.path.join(exp_name_dir, 'parameters.json')
    assert os.path.isfile(json_path), (
        "No json configuration file found at {}".format(json_path))
    with open(json_path) as data_file:
        params = json.load(data_file)
    """ train dataloader """
    data_loader = TrainDataLoader(args.train_path)
    """ 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))
    """ Model on gpu """
    model = TrackerSiamRPN(params)
    #model = model.cuda()
    cudnn.benchmark = True
    """ 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
    for epoch in range(start, args.max_epoches):
        #cur_lr = adjust_learning_rate(params["lr"], optimizer, epoch, gamma=0.1)
        index_list = range(data_loader.__len__())
        for example in tqdm(range(1000)):  # args.max_batches
            ret = data_loader.__get__(random.choice(index_list))

            closs, rloss, loss, reg_pred, reg_target, pos_index, neg_index, cur_lr = model.step(
                ret, epoch, backward=True)

            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())
            steps += 1

            if example % 1000 == 0:
                print(
                    "Epoch:{:04d}\texample:{:06d}/{:06d}({:.2f})%\tlr:{:.7f}\tcloss:{:.4f}\trloss:{:.4f}\ttloss:{:.4f}"
                    .format((epoch + 1), steps, args.max_batches,
                            100 * (steps) / args.max_batches, cur_lr,
                            closses.avg, rlosses.avg, tlosses.avg))
        """save model"""
        model_save_dir_pth = '{}/model'.format(exp_name_dir)
        if not os.path.exists(model_save_dir_pth):
            os.makedirs(model_save_dir_pth)
        net_path = os.path.join(model_save_dir_pth,
                                'model_e%d.pth' % (epoch + 1))
        torch.save(model.net.state_dict(), net_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))
    with tf.GradientTape() as tape:
        prediction = model(sample_batch_tf)
        
    label_batch = np.reshape(label_batch, newshape=prediction.shape)
    sample_loss = loss_object(label_batch, prediction)
    sample_loss = tf.convert_to_tensor(sample_loss)
    # Get the gradients of the loss w.r.t to the sample batch.
    gradient = tape.gradient(sample_loss, sample_batch_tf, unconnected_gradients='zero')

    # Get the sign of the gradients to create the perturbation
    signed_grad = [tf.sign(gradient[0]), tf.sign(gradient[1])]
    # signed_grad = gradient
    return signed_grad, prediction

train_data_loader = TrainDataLoader()
num_batches = train_data_loader.__len__()

for i in range(num_batches):
    print("Training on batch: {}".format(i))
    sample_batch, label_batch = train_data_loader[i]
    perturbation, prediction = create_adversarial_pattern(sample_batch, label_batch)
    adv_sample_batch = [sample_batch[0] + H.epsilon * perturbation[0], sample_batch[1] + H.epsilon * perturbation[1]]
    adv_model.train_on_batch(adv_sample_batch, label_batch)

perturbation, prediction = create_adversarial_pattern(sample_batch, label_batch, model=adv_model)
adv_sample_batch = [sample_batch[0] + H.epsilon * perturbation[0], sample_batch[1] + H.epsilon * perturbation[1]]
adv_prediction = model.predict(adv_sample_batch)

prediction = [0 if pred < 0.5 else 1 for pred in prediction]
adv_prediction = [0 if pred < 0.5 else 1 for pred in adv_prediction]
misclassification_rate = np.logical_xor(adv_prediction, prediction).sum() / sample_batch[0].shape[0]