# 第一帧的bbox
init_info["init_bbox"]=[420,136,53,130]
frames_path = "/root/stark/Lsot_TIR/airplane/img"
frames_list = ['{}/{:08d}.jpg'.format(frames_path,frame_numbe) for frame_number in range(1, len(os.listdir(frames_path)) + 1)]
image=cv.imread('/root/stark/Lsot_TIR/airplane/img/00000001.jpg')
image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
start_time = time.time()
out = tracker.initialize(image, init_info)

# 即tracker.initialize()
def initialize(self, image, info: dict):

    z_patch_arr, _, z_amask_arr = sample_target(image, info['init_bbox'], self.params.template_factor,
                                                output_sz=self.params.template_size)
    template, template_mask = self.preprocessor.process(z_patch_arr, z_amask_arr)
    #print(template.dtype)
    template= np.ascontiguousarray(template.astype('float32'))  # 将内存连续排列
    template = Tensor(template) # 将numpy转为转为Tensor类
    template_mask = np.ascontiguousarray(template_mask.astype('float32'))  # 将内存连续排列
    template_mask = Tensor(template_mask) # 将numpy转为转为Tensor类
    erro = np.zeros(1)
    erro = Tensor(np.ascontiguousarray(erro.astype('float32')))
    
           
    model_backbone_path = "/root/stark/backbone_bottleneck_pe.om"  # 模型路径
    # 模型推理
    model = base.model(modelPath=model_backbone_path, deviceId=0)  # 初始化 base.model 类
    print("model:",model)
    print("输入：",model.input_shape(0),model.input_shape(1),model.input_shape(2))
    print("输出：",model.output_shape(0),model.output_shape(1),model.output_shape(2))
    
    
    self.ort_outs_z = model.infer(template,template_mask,erro)
    print("推理后的结果：",self.ort_outs_z)
    #print(template, template_mask)
    # forward the template once
    #ort_inputs = {'img_z': template, 'mask_z': template_mask}
    #self.ort_outs_z = self.ort_sess_z.run(None, ort_inputs)
    #model_complete_path = "/root/stark/complete.om"  # 模型路径
    
    # save states
    self.state = info['init_bbox']
    self.frame_id = 0

# sample_target()
def sample_target(im, target_bb, search_area_factor, output_sz=None, mask=None):
    """ Extracts a square crop centered at target_bb box, of area search_area_factor^2 times target_bb area

    args:
        im - cv image
        target_bb - target box [x, y, w, h]
        search_area_factor - Ratio of crop size to target size
        output_sz - (float) Size to which the extracted crop is resized (always square). If None, no resizing is done.

    returns:
        cv image - extracted crop
        float - the factor by which the crop has been resized to make the crop size equal output_size
    """
    if not isinstance(target_bb, list):
        x, y, w, h = target_bb.tolist()
    else:
        x, y, w, h = target_bb
    # Crop image
    crop_sz = math.ceil(math.sqrt(w * h) * search_area_factor)

    if crop_sz < 1:
        raise Exception('Too small bounding box.')

    x1 = int(round(x + 0.5 * w - crop_sz * 0.5))
    x2 = int(x1 + crop_sz)

    y1 = int(round(y + 0.5 * h - crop_sz * 0.5))
    y2 = int(y1 + crop_sz)

    x1_pad = max(0, -x1)
    x2_pad = max(x2 - im.shape[1] + 1, 0)

    y1_pad = max(0, -y1)
    y2_pad = max(y2 - im.shape[0] + 1, 0)

    # Crop target
    im_crop = im[y1 + y1_pad:y2 - y2_pad, x1 + x1_pad:x2 - x2_pad, :]
    if mask is not None:
        mask_crop = mask[y1 + y1_pad:y2 - y2_pad, x1 + x1_pad:x2 - x2_pad]

    # Pad
    im_crop_padded = cv.copyMakeBorder(im_crop, y1_pad, y2_pad, x1_pad, x2_pad, cv.BORDER_CONSTANT)
    # deal with attention mask
    H, W, _ = im_crop_padded.shape
    att_mask = np.ones((H,W))
    end_x, end_y = -x2_pad, -y2_pad
    if y2_pad == 0:
        end_y = None
    if x2_pad == 0:
        end_x = None
    att_mask[y1_pad:end_y, x1_pad:end_x] = 0
    if mask is not None:
        mask_crop_padded = F.pad(mask_crop, pad=(x1_pad, x2_pad, y1_pad, y2_pad), mode='constant', value=0)

    if output_sz is not None:
        resize_factor = output_sz / crop_sz
        im_crop_padded = cv.resize(im_crop_padded, (output_sz, output_sz))
        att_mask = cv.resize(att_mask, (output_sz, output_sz)).astype(np.bool_)
        if mask is None:
            return im_crop_padded, resize_factor, att_mask
        mask_crop_padded = \
        F.interpolate(mask_crop_padded[None, None], (output_sz, output_sz), mode='bilinear', align_corners=False)[0, 0]
        return im_crop_padded, resize_factor, att_mask, mask_crop_padded

    else:
        if mask is None:
            return im_crop_padded, att_mask.astype(np.bool_), 1.0
        return im_crop_padded, 1.0, att_mask.astype(np.bool_), mask_crop_padded

# self.preprocessor.process()
class PreprocessorX_onnx(object):
    def __init__(self):
        self.mean = np.array([0.485, 0.456, 0.406]).reshape((1, 3, 1, 1))
        self.std = np.array([0.229, 0.224, 0.225]).reshape((1, 3, 1, 1))

    def process(self, img_arr: np.ndarray, amask_arr: np.ndarray):
        """img_arr: (H,W,3), amask_arr: (H,W)"""
        # Deal with the image patch
        img_arr_4d = img_arr[np.newaxis, :, :, :].transpose(0, 3, 1, 2)
        img_arr_4d = (img_arr_4d / 255.0 - self.mean) / self.std  # (1, 3, H, W)
        # Deal with the attention mask
        amask_arr_3d = amask_arr[np.newaxis, :, :]  # (1,H,W)
        return img_arr_4d.astype(np.float32), amask_arr_3d.astype(np.bool)