使用作者的预训练权重在自己的数据集上finetune Siamese Rpn

使用作者的预训练权重在自己的数据集上finetune Siamese Rpn，以获得更好的跟踪效果

完整工程：https://github.com/woshiwwwppp/Finetune-Simese-RPN-

训练主程序train.py

# -*- coding: utf-8 -*-
import os
import random
import sys; sys.path.append('../')
import torch
import torch.nn as nn
import numpy as np
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
import argparse
from run_SiamRPN import TrainDataLoader
from shapely.geometry import Polygon
from tensorboardX import SummaryWriter
from os.path import realpath, dirname, join
from net import SiamRPNvot

parser = argparse.ArgumentParser(description='PyTorch SiameseRPN Training')

parser.add_argument('--train_path', default='D:\\uav_frame\\00', metavar='DIR',help='path to dataset')

parser.add_argument('--weight_dir', default='D:\\project\\py\\yolo3\\Siamese-RPN-pytorch-master', metavar='DIR',help='path to weight')

parser.add_argument('--checkpoint_path', default=None, help='resume')

parser.add_argument('--max_epoches', default=10000, type=int, metavar='N', help='number of total epochs to run')

parser.add_argument('--max_batches', default=0, type=int, metavar='N', help='number of batch in one epoch')

parser.add_argument('--init_type',  default='xavier', type=str, metavar='INIT', help='init net')

parser.add_argument('--lr', default=0.0005, type=float, metavar='LR', help='initial learning rate')

parser.add_argument('--momentum', default=0.9, type=float, metavar='momentum', help='momentum')

parser.add_argument('--weight_decay', '--wd', default=5e-5, type=float, metavar='W', help='weight decay (default: 1e-4)')

parser.add_argument('--debug', default=False, type=bool,  help='whether to debug')

def main():
    args = parser.parse_args()
    """ compute max_batches """
    for root, dirs, files in os.walk(args.train_path):
        for dirnames in dirs:
            dir_path = os.path.join(root, dirnames)
            args.max_batches += len(os.listdir(dir_path))

    """ Model on gpu """
    model = SiamRPNvot()
    model = model.cuda()
    model.load_state_dict(torch.load(join(realpath(dirname(__file__)), 'SiamRPNVOT.model')))
    model.train().cuda()
    cudnn.benchmark = True

    """ train dataloader """
    data_loader = TrainDataLoader(args.train_path,model)
    if not os.path.exists(args.weight_dir):
        os.makedirs(args.weight_dir)


    """ 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
    writer = SummaryWriter()
    for epoch in range(args.max_epoches):
        cur_lr = adjust_learning_rate(args.lr, optimizer, epoch, gamma=0.1)
        index_list = range(data_loader.__len__())#获取数据集的长度
        losss=[0.0,0.0,0.0]

        for example in range(args.max_batches):
            ret = data_loader.__get__(random.choice(index_list))
            template = ret['temple'].cuda()
            detection= ret['detection'].cuda()
            pos_neg_diff = ret['pos_neg_diff_tensor'].cuda()
            model.temple(template)
            rout,cout = model(detection)
            cout = cout.squeeze().permute(1, 2, 0).reshape(-1, 2)
            rout = rout.squeeze().permute(1, 2, 0).reshape(-1, 4)

            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
            losss[0]=closses.avg
            losss[1] = rlosses.avg
            losss[2] = tlosses.avg
        print("Epoch:{:04d}\tcloss:{:.4f}\trloss:{:.4f}\ttloss:{:.4f}".format(epoch,  closses.avg, rlosses.avg, tlosses.avg ))
        writer.add_scalar("closses", losss[0], epoch)
        writer.add_scalar("rlosses", losss[1], epoch)
        writer.add_scalar("tlosses", losss[2], epoch)
        if steps % 150 == 0:
            file_path = os.path.join(args.weight_dir, 'weights-{:07d}.pth'.format(steps))
            state = {
            'epoch' :epoch+1,
            'state_dict' :model.state_dict(),
            'optimizer' : optimizer.state_dict(),
            }
            torch.save(state, file_path)

def intersection(g, p):
    g = Polygon(g[:8].reshape((4, 2)))
    p = Polygon(p[:8].reshape((4, 2)))
    if not g.is_valid or not p.is_valid:
        return 0
    inter = Polygon(g).intersection(Polygon(p)).area
    union = g.area + p.area - inter
    if union == 0:
        return 0
    else:
        return inter/union

def standard_nms(S, thres):
    """ use pre_thres to filter """
    index = np.where(S[:, 8] > thres)[0]
    S = S[index] # ~ 100, 4

    # Then use standard nms
    order = np.argsort(S[:, 8])[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        ovr = np.array([intersection(S[i], S[t]) for t in order[1:]])

        inds = np.where(ovr <= thres)[0]
        order = order[inds+1]
    return S[keep]

def reshape(x):
    t = np.array(x, dtype = np.float32)
    return t.reshape(-1, 1)

class MultiBoxLoss(nn.Module):
    def __init__(self):
        super(MultiBoxLoss, self).__init__()

    def forward(self, predictions, targets):
        #print('+++++++++++++++++++++++++++++++++++++++++++++++++++')
        cout, rout = predictions
        """ class """

        class_pred, class_target = cout, targets[:, 0].long()
        # pos_index , neg_index    = list(np.where(class_target == 1)[0]), list(np.where(class_target == 0)[0])
        #获取正负样本的位置
        pos_index, neg_index = list(np.where( class_target.cuda().data.cpu().numpy() == 1)[0]), list(np.where( class_target.cuda().data.cpu().numpy() == 0)[0])
        pos_num, neg_num         = len(pos_index), len(neg_index)
        #列表相加是拼起来再索引
        class_pred, class_target = class_pred[pos_index + neg_index], class_target[pos_index + neg_index]

        closs = F.cross_entropy(class_pred, class_target, size_average=False, reduce=False)
        closs = torch.div(torch.sum(closs), 64)

        """ regression """
        reg_pred = rout
        reg_target = targets[:, 1:]
        rloss = F.smooth_l1_loss(reg_pred, reg_target, size_average=False, reduce=False) #1445, 4
        rloss = torch.div(torch.sum(rloss, dim = 1), 4)
        rloss = torch.div(torch.sum(rloss[pos_index]), 16)

        loss = closs + rloss
        return closs, rloss, loss, reg_pred, reg_target, pos_index, neg_index

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def adjust_learning_rate(lr, optimizer, epoch, gamma=0.1):
    """Sets the learning rate to the initial LR decayed 0.9 every 50 epochs"""
    lr = lr * (0.9 ** (epoch // 1))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr

if __name__ == '__main__':
    main()

 

数据读取和标签制作

# --------------------------------------------------------
# DaSiamRPN
# Licensed under The MIT License
# Written by Qiang Wang (wangqiang2015 at ia.ac.cn)
# --------------------------------------------------------
import numpy as np
from torch.autograd import Variable
import torch
import torch.nn.functional as F
import random
import os
import os.path as osp
from utils import get_subwindow_tracking
import xml.etree.ElementTree as ET
import cv2
import sys
from net import SiamRPNvot
def generate_anchor(total_stride, scales, ratios, score_size):
    anchor_num = len(ratios) * len(scales)
    anchor = np.zeros((anchor_num, 4),  dtype=np.float32)
    size = total_stride * total_stride
    count = 0
    for ratio in ratios:
        # ws = int(np.sqrt(size * 1.0 / ratio))
        ws = int(np.sqrt(size / ratio))
        hs = int(ws * ratio)
        for scale in scales:
            wws = ws * scale
            hhs = hs * scale
            anchor[count, 0] = 0
            anchor[count, 1] = 0
            anchor[count, 2] = wws
            anchor[count, 3] = hhs
            count += 1

    anchor = np.tile(anchor, score_size * score_size).reshape((-1, 4))
    ori = - (score_size / 2) * total_stride
    xx, yy = np.meshgrid([ori + total_stride * dx for dx in range(score_size)],
                         [ori + total_stride * dy for dy in range(score_size)])
    xx, yy = np.tile(xx.flatten(), (anchor_num, 1)).flatten(), \
             np.tile(yy.flatten(), (anchor_num, 1)).flatten()
    anchor[:, 0], anchor[:, 1] = xx.astype(np.float32), yy.astype(np.float32)
    return anchor

class TrackerConfig(object):
    # These are the default hyper-params for DaSiamRPN 0.3827
    windowing = 'cosine'  # to penalize large displacements [cosine/uniform]
    # Params from the network architecture, have to be consistent with the training
    exemplar_size = 127  # input z size
    instance_size = 271  # input x size (search region)
    total_stride = 8
    score_size = (instance_size-exemplar_size)/total_stride+1
    context_amount = 0.5  # context amount for the exemplar
    ratios = [0.33, 0.5, 1, 2, 3]
    scales = [8, ]
    anchor_num = len(ratios) * len(scales)
    anchor = []
    penalty_k = 0.055
    window_influence = 0.42
    lr = 0.295
    # adaptive change search region #
    adaptive = True

    def update(self, cfg):
        for k, v in cfg.items():
            setattr(self, k, v)
        self.score_size = (self.instance_size - self.exemplar_size) / self.total_stride + 1

def tracker_eval(net, x_crop, target_pos, target_sz, window, scale_z, p):
    delta, score = net(x_crop)

    delta = delta.permute(1, 2, 3, 0).contiguous().view(4, -1).data.cpu().numpy()
    score = F.softmax(score.permute(1, 2, 3, 0).contiguous().view(2, -1), dim=0).data[1, :].cpu().numpy()

    delta[0, :] = delta[0, :] * p.anchor[:, 2] + p.anchor[:, 0]
    delta[1, :] = delta[1, :] * p.anchor[:, 3] + p.anchor[:, 1]
    delta[2, :] = np.exp(delta[2, :]) * p.anchor[:, 2]
    delta[3, :] = np.exp(delta[3, :]) * p.anchor[:, 3]

    def change(r):
        return np.maximum(r, 1./r)

    def sz(w, h):
        pad = (w + h) * 0.5
        sz2 = (w + pad) * (h + pad)
        return np.sqrt(sz2)

    def sz_wh(wh):
        pad = (wh[0] + wh[1]) * 0.5
        sz2 = (wh[0] + pad) * (wh[1] + pad)
        return np.sqrt(sz2)

    # size penalty
    s_c = change(sz(delta[2, :], delta[3, :]) / (sz_wh(target_sz)))  # scale penalty
    r_c = change((target_sz[0] / target_sz[1]) / (delta[2, :] / delta[3, :]))  # ratio penalty

    penalty = np.exp(-(r_c * s_c - 1.) * p.penalty_k)
    pscore = penalty * score

    # window float
    pscore = pscore * (1 - p.window_influence) + window * p.window_influence
    best_pscore_id = np.argmax(pscore)

    target = delta[:, best_pscore_id] / scale_z
    target_sz = target_sz / scale_z
    lr = penalty[best_pscore_id] * score[best_pscore_id] * p.lr

    res_x = target[0] + target_pos[0]
    res_y = target[1] + target_pos[1]

    res_w = target_sz[0] * (1 - lr) + target[2] * lr
    res_h = target_sz[1] * (1 - lr) + target[3] * lr

    target_pos = np.array([res_x, res_y])
    target_sz = np.array([res_w, res_h])
    return target_pos, target_sz, score[best_pscore_id]

def SiamRPN_init(im, target_pos, target_sz, net):
    state = dict()
    p = TrackerConfig()
    p.update(net.cfg)
    state['im_h'] = im.shape[0]
    state['im_w'] = im.shape[1]

    if p.adaptive:
        if ((target_sz[0] * target_sz[1]) / float(state['im_h'] * state['im_w'])) < 0.004:
            p.instance_size = 287  # small object big search region
        else:
            p.instance_size = 271

        p.score_size = (p.instance_size - p.exemplar_size) / p.total_stride + 1

    p.anchor = generate_anchor(p.total_stride, p.scales, p.ratios, int(p.score_size))

    avg_chans = np.mean(im, axis=(0, 1))#图像均值

    wc_z = target_sz[0] + p.context_amount * sum(target_sz)
    hc_z = target_sz[1] + p.context_amount * sum(target_sz)
    s_z = round(np.sqrt(wc_z * hc_z))
    # initialize the exemplar
    z_crop = get_subwindow_tracking(im, target_pos, p.exemplar_size, s_z, avg_chans)

    z = Variable(z_crop.unsqueeze(0))
    net.temple(z.cuda())

    if p.windowing == 'cosine':
        window = np.outer(np.hanning(p.score_size), np.hanning(p.score_size))
    elif p.windowing == 'uniform':
        window = np.ones((p.score_size, p.score_size))
    window = np.tile(window.flatten(), p.anchor_num)

    state['p'] = p
    state['net'] = net
    state['avg_chans'] = avg_chans
    state['window'] = window
    state['target_pos'] = target_pos
    state['target_sz'] = target_sz
    return state

def SiamRPN_track(state, im):
    p = state['p']
    net = state['net']
    avg_chans = state['avg_chans']
    window = state['window']
    target_pos = state['target_pos']
    target_sz = state['target_sz']

    wc_z = target_sz[1] + p.context_amount * sum(target_sz)
    hc_z = target_sz[0] + p.context_amount * sum(target_sz)
    s_z = np.sqrt(wc_z * hc_z)
    scale_z = p.exemplar_size / s_z
    d_search = (p.instance_size - p.exemplar_size) / 2
    pad = d_search / scale_z
    s_x = s_z + 2 * pad

    # extract scaled crops for search region x at previous target position
    x_crop = Variable(get_subwindow_tracking(im, target_pos, p.instance_size, round(s_x), avg_chans).unsqueeze(0))

    target_pos, target_sz, score = tracker_eval(net, x_crop.cuda(), target_pos, target_sz * scale_z, window, scale_z, p)
    target_pos[0] = max(0, min(state['im_w'], target_pos[0]))
    target_pos[1] = max(0, min(state['im_h'], target_pos[1]))
    target_sz[0] = max(10, min(state['im_w'], target_sz[0]))
    target_sz[1] = max(10, min(state['im_h'], target_sz[1]))
    state['target_pos'] = target_pos
    state['target_sz'] = target_sz
    state['score'] = score
    return state

class Anchor_ms(object):
    """
    stable version for anchor generator
    """
    def __init__(self, feature_w, feature_h):
        self.w = feature_w
        self.h = feature_h
        self.base = 64  # base size for anchor box
        self.stride = 15  # center point shift stride
        self.scale = [1 / 3, 1 / 2, 1, 2, 3]  # aspect ratio
        self.anchors = self.gen_anchors()  # xywh
        self.eps = 0.01

    def gen_single_anchor(self):
        scale = np.array(self.scale, dtype=np.float32)
        s = self.base * self.base
        w, h = np.sqrt(s / scale), np.sqrt(s * scale)
        c_x, c_y = (self.stride - 1) // 2, (self.stride - 1) // 2
        anchor = np.vstack([c_x * np.ones_like(scale, dtype=np.float32), c_y * np.ones_like(scale, dtype=np.float32), w,
                            h]).transpose()
        anchor = self.center_to_corner(anchor)
        return anchor

    def gen_anchors(self):
        anchor = self.gen_single_anchor()
        k = anchor.shape[0]
        delta_x, delta_y = [x * self.stride for x in range(self.w)], [y * self.stride for y in range(self.h)]
        shift_x, shift_y = np.meshgrid(delta_x, delta_y)
        shifts = np.vstack([shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel()]).transpose()
        a = shifts.shape[0]
        anchors = (anchor.reshape((1, k, 4)) + shifts.reshape((a, 1, 4))).reshape((a * k, 4))  # corner format
        anchors = self.corner_to_center(anchors)
        return anchors

    # float
    def diff_anchor_gt(self, gt):
        eps = self.eps
        anchors, gt = self.anchors.copy(), gt.copy()
        diff = np.zeros_like(anchors, dtype=np.float32)  # 生成和anchor一样大小的0矩阵
        # detection帧的groundtruth，四个坐标值分别反变换后得到rpn滑框的dx,dy,dw,dh
        diff[:, 0] = (gt[0] - anchors[:, 0]) / (anchors[:, 2] + eps)  #
        diff[:, 1] = (gt[1] - anchors[:, 1]) / (anchors[:, 3] + eps)
        diff[:, 2] = np.log((gt[2] + eps) / (anchors[:, 2] + eps))
        diff[:, 3] = np.log((gt[3] + eps) / (anchors[:, 3] + eps))
        return diff

    # float
    def center_to_corner(self, box):
        box = box.copy()
        box_ = np.zeros_like(box, dtype=np.float32)
        box_[:, 0] = box[:, 0] - (box[:, 2] - 1) / 2
        box_[:, 1] = box[:, 1] - (box[:, 3] - 1) / 2
        box_[:, 2] = box[:, 0] + (box[:, 2] - 1) / 2
        box_[:, 3] = box[:, 1] + (box[:, 3] - 1) / 2
        box_ = box_.astype(np.float32)
        return box_

    # float
    def corner_to_center(self, box):
        box = box.copy()
        box_ = np.zeros_like(box, dtype=np.float32)
        box_[:, 0] = box[:, 0] + (box[:, 2] - box[:, 0]) / 2
        box_[:, 1] = box[:, 1] + (box[:, 3] - box[:, 1]) / 2
        box_[:, 2] = (box[:, 2] - box[:, 0])
        box_[:, 3] = (box[:, 3] - box[:, 1])
        box_ = box_.astype(np.float32)
        return box_

    def pos_neg_anchor(self, gt, pos_num=16, neg_num=48, threshold_pos=0.5, threshold_neg=0.1):
        #根据detection帧的box与自动生成的box计算iou，得出正负样本
        gt = gt.copy()
        gt_corner = self.center_to_corner(np.array(gt, dtype=np.float32).reshape(1, 4))
        an_corner = self.center_to_corner(np.array(self.anchors, dtype=np.float32))
        # 计算groundtruth与生成的1445个框的iou
        iou_value = self.iou(an_corner, gt_corner).reshape(-1)  # (1445)
        max_iou = max(iou_value)
        pos, neg = np.zeros_like(iou_value, dtype=np.int32), np.zeros_like(iou_value, dtype=np.int32)

        # pos
        # 选取得分最高的30个框
        pos_cand = np.argsort(iou_value)[::-1][:30]
        # 从这30个框中选择16个作为正样本
        pos_index = np.random.choice(pos_cand, pos_num, replace=False)
        if max_iou > threshold_pos:
            pos[pos_index] = 1  # 对应pos_index的pos置1

        # neg
        neg_cand = np.where(iou_value < threshold_neg)[0]  # 选择iou小于threshold_neg的为负样本
        neg_ind = np.random.choice(neg_cand, neg_num, replace=False)  # 随机选择48个负样本
        neg[neg_ind] = 1

        return pos, neg

    def iou(self, box1, box2):
        box1, box2 = box1.copy(), box2.copy()
        N = box1.shape[0]
        K = box2.shape[0]
        box1 = np.array(box1.reshape((N, 1, 4))) + np.zeros((1, K, 4))  # box1=[N,K,4]
        box2 = np.array(box2.reshape((1, K, 4))) + np.zeros((N, 1, 4))  # box1=[N,K,4]
        x_max = np.max(np.stack((box1[:, :, 0], box2[:, :, 0]), axis=-1), axis=2)
        x_min = np.min(np.stack((box1[:, :, 2], box2[:, :, 2]), axis=-1), axis=2)
        y_max = np.max(np.stack((box1[:, :, 1], box2[:, :, 1]), axis=-1), axis=2)
        y_min = np.min(np.stack((box1[:, :, 3], box2[:, :, 3]), axis=-1), axis=2)
        tb = x_min - x_max
        lr = y_min - y_max
        tb[np.where(tb < 0)] = 0
        lr[np.where(lr < 0)] = 0
        over_square = tb * lr
        all_square = (box1[:, :, 2] - box1[:, :, 0]) * (box1[:, :, 3] - box1[:, :, 1]) + (
                    box2[:, :, 2] - box2[:, :, 0]) * (box2[:, :, 3] - box2[:, :, 1]) - over_square
        return over_square / all_square

class TrainDataLoader(object):

    # out_feature最后滑框和类别特征图的大小
    def __init__(self, img_dir_path, net,out_feature=19, max_inter=80, check=False ,tmp_dir='../tmp/visualization'):  #
        assert osp.isdir(img_dir_path), 'input img_dir_path error'

        self.img_dir_path = img_dir_path  # this is a root dir contain subclass
        self.max_inter = max_inter
        self.sub_class_dir = [sub_class_dir for sub_class_dir in os.listdir(img_dir_path) if
                              os.path.isdir(os.path.join(img_dir_path, sub_class_dir))]
        self.anchor_generator = Anchor_ms(out_feature, out_feature)
        self.anchors = self.anchor_generator.gen_anchors()  # centor，依据17*17的特征图，每个点生成5个框，17*17*5=1445
        self.ret = {}
        self.check = check
        self.tmp_dir = self.init_dir(tmp_dir)
        self.ret['tmp_dir'] = tmp_dir
        self.ret['check'] = check
        self.count = 0
        self.ret['p']=TrackerConfig()
        self.ret['p'].update(net.cfg)

    def init_dir(self, tmp_dir):
        if not osp.exists(tmp_dir):
            os.makedirs(tmp_dir)
        return tmp_dir

    # def get_transform_for_train(self):
    #     transform_list = []
    #     transform_list.append(transforms.ToTensor())
    #     transform_list.append(transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
    #     return transforms.Compose(transform_list)
    #
    # # tuple  计算图像三个通道的均值
    # def _average(self):
    #     assert self.ret.__contains__('template_img_path'), 'no template path'
    #     assert self.ret.__contains__('detection_img_path'), 'no detection path'
    #     template = Image.open(self.ret['template_img_path'])
    #     detection = Image.open(self.ret['detection_img_path'])
    #
    #     mean_template = tuple(map(round, ImageStat.Stat(template).mean))
    #     mean_detection = tuple(map(round, ImageStat.Stat(detection).mean))
    #     self.ret['mean_template'] = mean_template
    #     self.ret['mean_detection'] = mean_detection

    def _pick_img_pairs(self, index_of_subclass):
        # img_dir_path -> sub_class_dir_path -> template_img_path
        # use index_of_subclass to select a sub directory
        assert index_of_subclass < len(self.sub_class_dir), 'index_of_subclass should less than total classes'
        sub_class_dir_basename = self.sub_class_dir[index_of_subclass]
        sub_class_dir_path = os.path.join(self.img_dir_path, sub_class_dir_basename)
        sub_class_img_name = [img_name for img_name in os.listdir(sub_class_dir_path) if
                              not img_name.find('.jpg') == -1]
        sub_class_img_name = sorted(sub_class_img_name)
        sub_class_img_num = len(sub_class_img_name)
        # select template, detection
        # ++++++++++++++++++++++++++++ add break in sequeence [0,0,0,0] ++++++++++++++++++++++++++++++++++
        if self.max_inter >= sub_class_img_num - 1:
            self.max_inter = sub_class_img_num // 2
        template_index = np.clip(random.choice(range(0, max(1, sub_class_img_num - self.max_inter))), 0,
                                 sub_class_img_num - 1)  # 在子文件夹中随机选了一个template
        detection_index = np.clip(random.choice(range(1, max(2, self.max_inter))) + template_index, 0,
                                  sub_class_img_num - 1)  # 在子文件夹中随机选了一个template
        template_name, detection_name = sub_class_img_name[template_index], sub_class_img_name[detection_index]
        template_img_path, detection_img_path = osp.join(sub_class_dir_path, template_name), osp.join(
            sub_class_dir_path, detection_name)
        im_template = cv2.imread(template_img_path)
        im_detection = cv2.imread(detection_img_path)
        template_xml_path=template_img_path.split('.')[0]+'.xml'
        detection_xml_path = detection_img_path.split('.')[0] + '.xml'
        box_template=self.get_xywh_from_xml(template_xml_path)
        box_detection = self.get_xywh_from_xml(detection_xml_path)

        # load infomation of template and detection
        self.ret['img_template'] = im_template
        self.ret['img_detection'] = im_detection
        self.ret['template_target_pos'] = np.array([box_template[0],box_template[1]])
        self.ret['template_target_sz'] = np.array([box_template[2],box_template[3]])
        self.ret['detection_target_pos'] = np.array([box_detection[0],box_detection[1]])
        self.ret['detection_target_sz'] = np.array([box_detection[2], box_detection[3]])
        self.ret['anchors'] = self.anchors
        # self._average()  # 计算图像均值

    def get_xywh_from_xml(self,file):
        in_file = open(file)
        tree = ET.parse(in_file)
        root = tree.getroot()
        x,y,w,h=0,0,0,0
        for object in root.iter('object'):
            bndbox=object.find('bndbox')
            xmin= int(bndbox.find('xmin').text)
            xmax = int(bndbox.find('xmax').text)
            ymin = int(bndbox.find('ymin').text)
            ymax = int(bndbox.find('ymax').text)
            x=int((xmin+xmax)/2)
            y=int((ymin+ymax)/2)
            w=int(xmax-xmin)
            h=int(ymax-ymin)
        return [x,y,w,h]

    def _pad_crop_resize_template(self):
        self.ret['im_h'] = self.ret['img_template'].shape[0]
        self.ret['im_w'] = self.ret['img_template'].shape[1]
        self.ret['p'].score_size = (self.ret['p'].instance_size - self.ret['p'].exemplar_size) / self.ret['p'].total_stride + 1
        self.ret['p'].anchor = generate_anchor(self.ret['p'].total_stride, self.ret['p'].scales, self.ret['p'].ratios, int(self.ret['p'].score_size))
        avg_chans = np.mean(self.ret['img_template'], axis=(0, 1))  # 图像均值
        wc_z = self.ret['template_target_sz'][0] + self.ret['p'].context_amount * sum(self.ret['template_target_sz'])
        hc_z = self.ret['template_target_sz'][1] + self.ret['p'].context_amount * sum(self.ret['template_target_sz'])
        s_z = round(np.sqrt(wc_z * hc_z))
        # initialize the exemplar
        z_crop = get_subwindow_tracking(self.ret['img_template'], self.ret['template_target_pos'], self.ret['p'].exemplar_size, s_z, avg_chans)
        z = Variable(z_crop.unsqueeze(0))
        # net.temple(z.cuda())
        if self.ret['p'].windowing == 'cosine':
            window = np.outer(np.hanning(self.ret['p'].score_size), np.hanning(self.ret['p'].score_size))
        elif self.ret['p'].windowing == 'uniform':
            window = np.ones((self.ret['p'].score_size, self.ret['p'].score_size))
        window = np.tile(window.flatten(), self.ret['p'].anchor_num)
        self.ret['temple'] = z
        self.ret['avg_chans'] = avg_chans
        self.ret['window'] = window

    def _pad_crop_resize_detection(self):
        wc_z = self.ret['detection_target_sz'][1] + self.ret['p'].context_amount * sum(self.ret['detection_target_sz'])
        hc_z = self.ret['detection_target_sz'][0] + self.ret['p'].context_amount * sum(self.ret['detection_target_sz'])
        s_z = np.sqrt(wc_z * hc_z)
        scale_z = self.ret['p'].exemplar_size / s_z
        d_search = (self.ret['p'].instance_size - self.ret['p'].exemplar_size) / 2
        pad = d_search / scale_z
        s_x = s_z + 2 * pad
        avg_chans = np.mean(self.ret['img_detection'], axis=(0, 1))  # 图像均值
        # extract scaled crops for search region x at previous target position
        x_crop = Variable(get_subwindow_tracking(self.ret['img_detection'], self.ret['detection_target_pos'], self.ret['p'].instance_size, round(s_x), avg_chans).unsqueeze(0))
        self.ret['detection'] = x_crop


    def _generate_pos_neg_diff(self):

        # np.array((self.ret['detection_target_pos'][0],self.ret['detection_target_pos'][1], self.ret['detection_target_sz'][0], self.ret['detection_target_sz'][1]), dtype=np.int32)
        gt_box_in_detection = np.array((self.ret['detection_target_pos'][0],self.ret['detection_target_pos'][1], self.ret['detection_target_sz'][0], self.ret['detection_target_sz'][1]), dtype=np.int32)
        pos, neg = self.anchor_generator.pos_neg_anchor(gt_box_in_detection)  # 生成框的正负样本标记向量
        diff = self.anchor_generator.diff_anchor_gt(gt_box_in_detection)  # 生成回归框的dx,dy,dw,dh
        pos, neg, diff = pos.reshape((-1, 1)), neg.reshape((-1, 1)), diff.reshape((-1, 4))
        class_target = np.array([-100.] * self.anchors.shape[0], np.int32)

        # pos   获取pos的真实长度，并生成标记向量
        pos_index = np.where(pos == 1)[0]
        pos_num = len(pos_index)
        self.ret['pos_anchors'] = np.array(self.ret['anchors'][pos_index, :],dtype=np.int32) if not pos_num == 0 else None
        if pos_num > 0:
            class_target[pos_index] = 1

        # neg   获取neg的真实长度，并生成标记向量
        neg_index = np.where(neg == 1)[0]
        neg_num = len(neg_index)
        class_target[neg_index] = 0

   

        class_logits = class_target.reshape(-1, 1)  # 生成框的类别logits的值用于分类，负样本的值为-100，正样本的值为0
        pos_neg_diff = np.hstack((class_logits, diff))  # 将类别和坐标合并
        self.ret['pos_neg_diff'] = pos_neg_diff
        self.ret['pos_neg_diff_tensor'] = torch.Tensor(pos_neg_diff)

        return pos_neg_diff

  
    def to_torch(self,ndarray):
        if type(ndarray).__module__ == 'numpy':
            return torch.from_numpy(ndarray)
        elif not torch.is_tensor(ndarray):
            raise ValueError("Cannot convert {} to torch tensor"
                             .format(type(ndarray)))
        return ndarray

    def im_to_torch(self,img):
        img = np.transpose(img, (2, 0, 1))  # C*H*W
        img = self.to_torch(img).float()
        return img

    def get_subwindow_tracking(self,im, pos, model_sz, original_sz, avg_chans, out_mode='torch', new=False):
        if isinstance(pos, float):
            pos = [pos, pos]
        sz = original_sz
        im_sz = im.shape
        c = (original_sz + 1) / 2
        context_xmin = round(pos[0] - c)  # floor(pos(2) - sz(2) / 2);
        context_xmax = context_xmin + sz - 1
        context_ymin = round(pos[1] - c)  # floor(pos(1) - sz(1) / 2);
        context_ymax = context_ymin + sz - 1
        left_pad = int(max(0., -context_xmin))
        top_pad = int(max(0., -context_ymin))
        right_pad = int(max(0., context_xmax - im_sz[1] + 1))
        bottom_pad = int(max(0., context_ymax - im_sz[0] + 1))

        context_xmin = context_xmin + left_pad
        context_xmax = context_xmax + left_pad
        context_ymin = context_ymin + top_pad
        context_ymax = context_ymax + top_pad

        # zzp: a more easy speed version
        r, c, k = im.shape
        if any([top_pad, bottom_pad, left_pad, right_pad]):
            te_im = np.zeros((r + top_pad + bottom_pad, c + left_pad + right_pad, k),
                             np.uint8)  # 0 is better than 1 initialization
            te_im[top_pad:top_pad + r, left_pad:left_pad + c, :] = im
            if top_pad:
                te_im[0:top_pad, left_pad:left_pad + c, :] = avg_chans
            if bottom_pad:
                te_im[r + top_pad:, left_pad:left_pad + c, :] = avg_chans
            if left_pad:
                te_im[:, 0:left_pad, :] = avg_chans
            if right_pad:
                te_im[:, c + left_pad:, :] = avg_chans
            im_patch_original = te_im[int(context_ymin):int(context_ymax + 1), int(context_xmin):int(context_xmax + 1),
                                :]
        else:
            im_patch_original = im[int(context_ymin):int(context_ymax + 1), int(context_xmin):int(context_xmax + 1), :]

        if not np.array_equal(model_sz, original_sz):
            im_patch = cv2.resize(im_patch_original, (model_sz, model_sz))  # zzp: use cv to get a better speed
        else:
            im_patch = im_patch_original

        return self.im_to_torch(im_patch) if out_mode in 'torch' else im_patch

    def __get__(self, index):
        self._pick_img_pairs(index)
        self._pad_crop_resize_template()
        self._pad_crop_resize_detection()
        self._generate_pos_neg_diff()  # 生成框的1445*5的张量，代表每个框的类别，dx,dy,dw,dh
        # self._tranform()  # PIL to Tensor
        self.count += 1
        return self.ret

    def __len__(self):
        return len(self.sub_class_dir)



if __name__ == '__main__':
    # we will do a test for dataloader
    net = SiamRPNvot()
    loader = TrainDataLoader('D:\\uav_frame\\00',net ,check = True)
    #print(loader.__len__())
    index_list = range(loader.__len__())
    for i in range(1000):
        ret = loader.__get__(random.choice(index_list))
        label = ret['pos_neg_diff'][:, 0].reshape(-1)
        pos_index = list(np.where(label == 1)[0])
        pos_num = len(pos_index)
        print(pos_index)
        print(pos_num)
        if pos_num != 0 and pos_num != 16:
            print(pos_num)
            sys.exit(0)
        print(i)