コード例 #1
0
ファイル: simulator.py プロジェクト: yueshen95/re3-pytorch
def measure_occlusion(bbox,
                      trackedObjects,
                      cropSize=CROP_SIZE,
                      cropPad=CROP_PAD):
    image = np.zeros((int(cropSize), int(cropSize)), dtype=np.bool)
    fullBBox = scale_bbox(bbox, cropPad)
    fullBBoxXYWH = xyxy_to_xywh(fullBBox)
    fullBBoxXYWH[[2, 3]] = np.maximum(fullBBoxXYWH[[2, 3]], 1)

    # Now do all objects
    for obj in trackedObjects:
        boxPos = obj.get_bbox()
        boxPosXYWH = xyxy_to_xywh(boxPos)
        if IOU(boxPos, fullBBox) < 0.001:
            continue
        cropCoords = np.zeros(4)
        cropCoords = np.clip(boxPos - fullBBox[[0, 1, 0, 1]], 0,
                             fullBBoxXYWH[[2, 3, 2, 3]])
        cropCoords *= cropSize * 1.0 / fullBBoxXYWH[[2, 3, 2, 3]]
        cropCoords = np.clip(np.round(cropCoords), 0, cropSize).astype(int)
        if (cropCoords[2] - cropCoords[0] < 1
                or cropCoords[3] - cropCoords[1] < 1):
            cropCoords[[0, 1]] = np.clip(cropCoords[[0, 1]] - 1, 0,
                                         cropSize).astype(int)
            cropCoords[[2, 3]] = np.clip(cropCoords[[2, 3]] + 1, 0,
                                         cropSize).astype(int)
        image[cropCoords[1]:cropCoords[3], cropCoords[0]:cropCoords[2]] = True
    return np.count_nonzero(image) * 1.0 / image.size
コード例 #2
0
def rectangle_guessing_loss(y_true, y_pred):
	y_true_np = K.eval(y_true)
	y_pred_np = K.eval(y_pred)
	y_true_iter = np.nditer(y_true_np, order='C')
	y_pred_iter = np.nditer(y_pred_np, order='C')
	loss = 0.0
	count = 0
	for y1, y2 in zip(y_true_iter, y_pred_iter):
		loss.append(IOU(y1, y2))
		count += 1
	
	return (loss / count)
コード例 #3
0
    def responsible_box(self, idx, label, predictions, sx, sy):
        max_iou = 0.0
        responsible = 0
        for b in range(self.B):
            boxA, boxB = self.get_real_boxes(label[5 * idx + 1:5 * idx + 5],
                                             predictions[5 * b + 1:5 * b + 5],
                                             sx, sy)
            iou = IOU(boxA, boxB)

            if max_iou < iou:
                max_iou = iou
                responsible = b
        return b, max_iou
コード例 #4
0
ファイル: dataset.py プロジェクト: V-9318/YoloNet
    def _kmeans(self, X):
        k = self.NAnchors
        print('kmeans with k = {0}'.format(k))
        x_num = X.shape[0]
        iterations = 0
        self.dataset['train'].append(ell)
        prev_assignments = np.ones(k)*(-1)
        iteration = 0
        old_distances = np.zeros((x_num, k))

        indices = [random.randrange(x_num) for i in range(k)]
        centroids = X[indices]
        anchor_dim = X.shape[1]

        while True:
            distances = []
            iteration += 1
            for i in range(x_num):
                d = 1 - IOU(ann_dims[i], centroids)
                distances.append(d)
            distances = np.array(distances)
            print("iteration {}: dists = {}".format(iteration, np.sum(np.abs(old_distances-distances))))

            #assign samples to centroids
            assignments = np.argmin(distances,axis=1)

            if (assignments == prev_assignments).all() :
                return centroids

            #calculate new centroids
            centroid_sums=np.zeros((k, anchor_dim), np.float)
            for i in range(x_num):
                centroid_sums[assignments[i]]+=X[i]
            for j in range(k):
                centroids[j] = centroid_sums[j]/(np.sum(assignments==j) + 1e-6)

            prev_assignments = assignments.copy()
            old_distances = distances.copy()
コード例 #5
0
def remove_bb(boxes):
    """ removes redundant bbs from list which have great iou with other bbs in list.
    """
    for i in range(len(boxes)):
        boxA = boxes[i]
        if boxA is None:
            continue
        boxA_rect = calc_rect(boxA)
        for j in range( len(boxes)):
            boxB = boxes[j]
            if boxB is None:
                continue
            if np.array_equal(boxA, boxB):
                continue
            boxB_rect = calc_rect(boxB)

            iou = IOU(boxA_rect, boxB_rect)
            if iou>0.3:
                boxes[j] =None
    #filter out None elements from list
    boxes_final = list(filter(None, boxes))

    return boxes_final
コード例 #6
0
ファイル: eval_lasot.py プロジェクト: shen338/tracker
            dscores = detection_result["instances"].scores.cpu().detach(
            ).numpy()
            center_pos = tracker.tracker.center_pos

            # print(tracker.tracker.center_pos, tracker.tracker.size)
            # print(dboxes, dscores)
            # print(x_crop.shape)
            all_outputs = tracker.track(frame, x_crop, scale_z,
                                        search_instance_size)
            # cv2.imwrite("test.jpg", frame)
            tboxes, tscores = all_outputs['bbox'], all_outputs['best_score']
            # print(tboxes, tscores)

            for idx, dbox in enumerate(dboxes):
                for tbox in tboxes:
                    if IOU(dbox, tbox) > 0.8:
                        dscores[idx] = 1

            # print(dscores)
            # Windows penalty for detection results

            input_size = int(search_instance_size / scale_z)
            hanning = np.hanning(input_size)
            window = np.outer(hanning, hanning)

            idx = 0
            for idx, dbox in enumerate(dboxes):
                x = int((dbox[0] + dbox[2]) / 2)
                y = int((dbox[1] + dbox[3]) / 2)
                # print(x, y, idx)
                dscores[idx] = dscores[idx] * (
コード例 #7
0
ファイル: train.py プロジェクト: alecGraves/unet-tf
            print('start testing {} samples...'.format(num_test_samples))
            for ti in range(num_test_samples // batch_size):
                x_batch, y_batch = next(test_generator)
                # tensorflow wants a different tensor order
                feed_dict = {
                    img: x_batch,
                    label: y_batch,
                }
                loss, pred_logits = sess.run([cross_entropy_loss, pred],
                                             feed_dict=feed_dict)
                pred_map_batch = np.argmax(pred_logits, axis=3)

        x_batch, y_batch = next(train_generator)
        feed_dict = {img: x_batch, label: y_batch}
        _, loss, summary, lr, pred_logits = sess.run([
            train_step, cross_entropy_loss, summary_merged, learning_rate, pred
        ],
                                                     feed_dict=feed_dict)
        global_step.assign(it).eval()
        train_writer.add_summary(summary, it)

        score = IOU(1 / (1 + np.exp(-pred_logits[0])), y_batch[0])

        if it % 20 == 0:
            try:
                print('[iter {}, epoch {}]: lr={} loss={}, IOU={}'.format(
                    it,
                    float(it) / steps_per_epoch, lr, loss, score))
            except:
                pass
コード例 #8
0
    idx += 1
    height, width, channel = img.shape

    neg_num = 0
    #先采样一定数量neg图片
    while neg_num < 50:

        #随机选取截取图像大小
        size = npr.randint(12, min(width, height) / 2)
        #随机选取左上坐标
        nx = npr.randint(0, width - size)
        ny = npr.randint(0, height - size)
        #截取box
        crop_box = np.array([nx, ny, nx + size, ny + size])
        #计算iou值
        Iou = IOU(crop_box, boxes)
        #截取图片并resize成12x12大小
        cropped_im = img[ny:ny + size, nx:nx + size, :]
        resized_im = cv2.resize(cropped_im, (12, 12),
                                interpolation=cv2.INTER_LINEAR)

        #iou值小于0.3判定为neg图像
        if np.max(Iou) < 0.3:
            save_file = os.path.join(neg_save_dir, '%s.jpg' % n_idx)
            f2.write(neg_save_dir + '/%s.jpg' % n_idx + ' 0\n')
            cv2.imwrite(save_file, resized_im)
            n_idx += 1
            neg_num += 1

    for box in boxes:
        #左上右下坐标
コード例 #9
0
    def box_head_evaluation(self, nms_boxes: list, nms_scores: list, nms_labels: list,
                            gt_boxes: list, gt_labels: list,
                            iou_thresh: float = 0.5) -> tuple:
        """
        Constructs matches list & scores list for every class

        Input
        -----
            nms_boxes: list:len(bz){(post_NMS_boxes_per_image,4)}  ([x1,y1,x2,y2] format)
            nms_scores: list:len(bz){(post_NMS_boxes_per_image)}   ( the score for the top class for
                                                                the regressed box)
            nms_labels: list:len(bz){(post_NMS_boxes_per_image)}   (top class of each regressed box)
            gt_bboxes: list: len(bz){(n_obj, 4)}([x1, y1, x2, y2] format)
            gt_labels: list: len(bz) {(n_obj)}

        Output
        -----
            matches:        (bz, post_NMS_boxes_per_image, 3)
            scores:         (bz, post_NMS_boxes_per_image, 3)
            num_true:       (3, )
            num_positives:  (3, )
        """

        matches = []
        scores = []
        num_trues = torch.zeros(1, 3)
        num_positives = torch.zeros(1, 3)

        batch_size = len(nms_labels)

        for bz in range(batch_size):
            match = torch.zeros(nms_labels[bz].shape[0], 3)
            score = torch.zeros(match.shape)
            num_true = torch.zeros(1, 3)
            num_positive = torch.zeros(num_true.shape)

            # calculate trues
            for obj_label in gt_labels[bz]:
                if obj_label > 0:
                    num_true[0, obj_label.type(torch.long) - 1] += 1

            # calculate positives
            for obj_label in nms_labels[bz]:
                num_positive[0, obj_label.type(torch.long)] += 1

            for i_pred, class_pred in enumerate(nms_labels[bz]):
                class_pred_scalar = class_pred.item()
                if class_pred_scalar + 1 in gt_labels[bz]:

                    # retrieve class gt label
                    i_gt_list = (gt_labels[bz] == class_pred + 1).nonzero(as_tuple=False).squeeze(0)

                    for i_gt in i_gt_list:
                        # retrieve masks
                        box_pred = nms_boxes[bz][i_pred]
                        box_gt = gt_boxes[bz][i_gt.item()]

                        # compute IOU
                        iou = IOU(box_pred, box_gt, mode='corner')

                        if iou > iou_thresh:
                            match[i_pred, class_pred] = 1

                # no matter how we always store the bbox scores
                score[i_pred, class_pred] = nms_scores[bz][i_pred]

            matches.append(match)
            scores.append(score)
            num_trues = torch.cat([num_trues, num_true], dim=0)
            num_positives = torch.cat([num_positives, num_positive], dim=0)

        return torch.cat(matches), torch.cat(scores), \
            torch.sum(num_trues, dim=0), torch.sum(num_positives, dim=0)
コード例 #10
0
def main(size):
    '''用于处理带有landmark的数据'''
    # size=args.input_size
    #是否对图像变换
    argument = True
    if size == 12:
        net = 'PNet'
    elif size == 24:
        net = 'RNet'
    elif size == 48:
        net = 'ONet'
    image_id = 0
    #数据输出路径
    OUTPUT = os.path.join(data_dir, str(size))
    if not os.path.exists(OUTPUT):
        os.mkdir(OUTPUT)
    #图片处理后输出路径
    dstdir = os.path.join(OUTPUT, 'train_%s_landmark_aug' % (net))
    if not os.path.exists(dstdir):
        os.mkdir(dstdir)
    #label记录txt
    ftxt = os.path.join(data_dir, 'trainImageList.txt')
    #记录label的txt
    f = open(os.path.join(OUTPUT, 'landmark_%d_aug.txt' % (size)), 'w')
    #获取图像路径,box,关键点
    data = getDataFromTxt(ftxt, data_dir)
    idx = 0
    for (imgPath, box, landmarkGt) in tqdm(data):
        #存储人脸图片和关键点
        F_imgs = []
        F_landmarks = []
        img = cv2.imread(imgPath)

        img_h, img_w, img_c = img.shape
        gt_box = np.array([box.left, box.top, box.right, box.bottom])
        #人脸图片
        f_face = img[box.top:box.bottom + 1, box.left:box.right + 1]
        #resize成网络输入大小
        f_face = cv2.resize(f_face, (size, size))

        landmark = np.zeros((5, 2))
        for index, one in enumerate(landmarkGt):
            #关键点相对于左上坐标偏移量并归一化
            rv = ((one[0] - gt_box[0]) / (gt_box[2] - gt_box[0]),
                  (one[1] - gt_box[1]) / (gt_box[3] - gt_box[1]))
            landmark[index] = rv
        F_imgs.append(f_face)
        F_landmarks.append(landmark.reshape(10))
        landmark = np.zeros((5, 2))
        if argument:
            #对图像变换
            idx = idx + 1
            x1, y1, x2, y2 = gt_box
            gt_w = x2 - x1 + 1
            gt_h = y2 - y1 + 1
            #除去过小图像
            if max(gt_w, gt_h) < 40 or x1 < 0 or y1 < 0:
                continue
            for i in range(10):
                #随机裁剪图像大小
                box_size = npr.randint(int(min(gt_w, gt_h) * 0.8),
                                       np.ceil(1.25 * max(gt_w, gt_h)))
                #随机左上坐标偏移量
                delta_x = npr.randint(-gt_w * 0.2, gt_w * 0.2)
                delta_y = npr.randint(-gt_h * 0.2, gt_h * 0.2)
                #计算左上坐标
                nx1 = int(max(x1 + gt_w / 2 - box_size / 2 + delta_x, 0))
                ny1 = int(max(y1 + gt_h / 2 - box_size / 2 + delta_y, 0))
                nx2 = nx1 + box_size
                ny2 = ny1 + box_size
                #除去超过边界的
                if nx2 > img_w or ny2 > img_h:
                    continue
                #裁剪边框,图片
                crop_box = np.array([nx1, ny1, nx2, ny2])
                cropped_im = img[ny1:ny2 + 1, nx1:nx2 + 1, :]
                resized_im = cv2.resize(cropped_im, (size, size))
                iou = IOU(crop_box, np.expand_dims(gt_box, 0))
                #只保留pos图像
                if iou > 0.65:
                    F_imgs.append(resized_im)
                    #关键点相对偏移
                    for index, one in enumerate(landmarkGt):
                        rv = ((one[0] - nx1) / box_size,
                              (one[1] - ny1) / box_size)
                        landmark[index] = rv
                    F_landmarks.append(landmark.reshape(10))
                    landmark = np.zeros((5, 2))
                    landmark_ = F_landmarks[-1].reshape(-1, 2)
                    box = BBox([nx1, ny1, nx2, ny2])
                    #镜像
                    if random.choice([0, 1]) > 0:
                        face_flipped, landmark_flipped = flip(
                            resized_im, landmark_)
                        face_flipped = cv2.resize(face_flipped, (size, size))
                        F_imgs.append(face_flipped)
                        F_landmarks.append(landmark_flipped.reshape(10))
                    #逆时针翻转
                    if random.choice([0, 1]) > 0:
                        face_rotated_by_alpha, landmark_rorated = rotate(
                            img, box, box.reprojectLandmark(landmark_), 5)
                        #关键点偏移
                        landmark_rorated = box.projectLandmark(
                            landmark_rorated)
                        face_rotated_by_alpha = cv2.resize(
                            face_rotated_by_alpha, (size, size))
                        F_imgs.append(face_rotated_by_alpha)
                        F_landmarks.append(landmark_rorated.reshape(10))

                        #左右翻转
                        face_flipped, landmark_flipped = flip(
                            face_rotated_by_alpha, landmark_rorated)
                        face_flipped = cv2.resize(face_flipped, (size, size))
                        F_imgs.append(face_flipped)
                        F_landmarks.append(landmark_flipped.reshape(10))
                    #顺时针翻转
                    if random.choice([0, 1]) > 0:
                        face_rotated_by_alpha, landmark_rorated = rotate(
                            img, box, box.reprojectLandmark(landmark_), -5)
                        #关键点偏移
                        landmark_rorated = box.projectLandmark(
                            landmark_rorated)
                        face_rotated_by_alpha = cv2.resize(
                            face_rotated_by_alpha, (size, size))
                        F_imgs.append(face_rotated_by_alpha)
                        F_landmarks.append(landmark_rorated.reshape(10))

                        #左右翻转
                        face_flipped, landmark_flipped = flip(
                            face_rotated_by_alpha, landmark_rorated)
                        face_flipped = cv2.resize(face_flipped, (size, size))
                        F_imgs.append(face_flipped)
                        F_landmarks.append(landmark_flipped.reshape(10))
        F_imgs, F_landmarks = np.asarray(F_imgs), np.asarray(F_landmarks)
        for i in range(len(F_imgs)):
            #剔除数据偏移量在[0,1]之间
            if np.sum(np.where(F_landmarks[i] <= 0, 1, 0)) > 0:
                continue
            if np.sum(np.where(F_landmarks[i] >= 1, 1, 0)) > 0:
                continue
            cv2.imwrite(os.path.join(dstdir, '%d.jpg' % (image_id)), F_imgs[i])
            landmarks = list(map(str, list(F_landmarks[i])))
            f.write(
                os.path.join(dstdir, '%d.jpg' % (image_id)) + ' -2 ' +
                ' '.join(landmarks) + '\n')
            image_id += 1
    f.close()
    return F_imgs, F_landmarks
コード例 #11
0
def compute_map(dataloader, checkpoint_file, device):
    # =========================== Pretrained ===============================
    # Put the path were you save the given pretrained model
    pretrained_path = '../pretrained/checkpoint680.pth'
    device = torch.device(
        'cuda') if torch.cuda.is_available() else torch.device('cpu')
    backbone, rpn = pretrained_models_680(pretrained_path)
    backbone = backbone.to(device)
    rpn = rpn.to(device)

    # ========================== Prep ============================
    metric_trackers = []
    for i in range(3):
        metric_trackers.append(MetricTracker(i))

    # ========================= Loading Model ==============================
    boxHead = BoxHead(Classes=3, P=7, device=device).to(device)
    if torch.cuda.is_available():
        checkpoint = torch.load(checkpoint_file)
    else:
        checkpoint = torch.load(checkpoint_file,
                                map_location=torch.device('cpu'))

    print("[INFO] Weight loaded from checkpoint file: {}".format(
        checkpoint_file))
    boxHead.load_state_dict(checkpoint['model_state_dict'])
    boxHead.eval()  # set to eval mode

    for iter, data in enumerate(tqdm(dataloader), 0):
        images = data['images'].to(device)
        assert len(images) == 1, "Only support batch_size == 1"
        # if (iter == 50):
        #     break

        labels_gt_all = data['labels'][0].to(device)
        bbox_gt_all = data["bbox"][0].to(device)
        index = data['index'][0]

        prob, clas, boxes = run_inference(images, index, backbone, rpn,
                                          boxHead)

        for tracker_i in range(3):
            # focus on one class
            target_class = tracker_i + 1

            labels_gt = labels_gt_all[labels_gt_all == target_class]
            bbox_gt = bbox_gt_all[labels_gt_all ==
                                  target_class]  #n,4   x,y,w,h

            clas_pred = clas[clas == target_class]  #   m
            prob_pred = prob[clas == target_class]  #   m
            boxes_pred = boxes[clas == target_class]  #  m,4  x1,y1,x2,y2

            boxes_pred_xywh = torch.zeros_like(boxes_pred,
                                               dtype=boxes_pred.dtype,
                                               device=boxes_pred.device)
            boxes_pred_xywh[:, 0] = (boxes_pred[:, 0] + boxes_pred[:, 2]) / 2
            boxes_pred_xywh[:, 1] = (boxes_pred[:, 1] + boxes_pred[:, 3]) / 2
            boxes_pred_xywh[:, 2] = boxes_pred[:, 2] - boxes_pred[:, 0]
            boxes_pred_xywh[:, 3] = boxes_pred[:, 3] - boxes_pred[:, 1]

            # determine if it is a match with bbox
            N_gt = len(bbox_gt)  #n
            N_pred = len(clas_pred)  #m

            tp_indicator = torch.zeros((N_pred, ))  #m
            match_indice = torch.zeros((N_pred, ))  #m
            if (N_pred != 0) and (N_gt != 0):
                # IOU matrix
                iou_mat = torch.zeros((N_pred, N_gt))  # m,n
                for x in range(N_pred):
                    for y in range(N_gt):
                        iou_mat[x, y] = IOU(
                            torch.unsqueeze(boxes_pred_xywh[x, :], 0),
                            torch.unsqueeze(bbox_gt[y, :], 0))
                    if torch.max(iou_mat[x, :]) >= 0.5:
                        tp_indicator[x] = 1
                        index = torch.argmax(iou_mat[x, :])
                        match_indice[x] = index

            metric_trackers[tracker_i].add_match(prob_pred, tp_indicator,
                                                 match_indice, N_gt)

    # compute map
    for i in range(3):
        metric_trackers[i].compute_precision_recall()
        recall, precision = metric_trackers[i].sorted_pr_curve()

        ap = metric_trackers[i].compute_ap()
        print("class_id: {}. ap: {}".format(i, ap))
        metric_trackers[i].reset()
コード例 #12
0
        # cv2.imwrite("test.jpg", np.squeeze(x_crop).transpose((1,2,0)))
        dboxes = detection_result["instances"].pred_boxes.tensor.cpu().detach().numpy()
        dscores = detection_result["instances"].scores.cpu().detach().numpy()
        center_pos = tracker.tracker.center_pos

        # print(tracker.tracker.center_pos, tracker.tracker.size)
        # print("detection results: ", len(dboxes), dscores)
        # print(x_crop.shape)
        all_outputs = tracker.track(frame, x_crop, scale_z, search_instance_size)
        # cv2.imwrite("test.jpg", frame)
        tboxes, tscores = all_outputs['bbox'], all_outputs['best_score']
        # print("tracking results: ", tboxes, tscores)

        for idx, dbox in enumerate(dboxes):
            for tbox in tboxes: 
                if IOU(dbox, tbox) > 0.8: 
                    dscores[idx] = 1   

        # print(dscores)
        # Windows penalty for detection results

        input_size = int(search_instance_size/scale_z)
        hanning = np.hanning(input_size)
        window = np.outer(hanning, hanning)

        idx = 0
        for idx, dbox in enumerate(dboxes):
            x = int((dbox[0] + dbox[2])/2)
            y = int((dbox[1] + dbox[3])/2)
            # print(x, y, idx)
            dscores[idx] = dscores[idx] * (1 - window_influence) + window[x, y]*window_influence
コード例 #13
0
    bboxes = np.array(bbox, dtype=np.float32).reshape(-1, 4)
    img = cv2.imread(os.path.join(img_dir, img_path + '.jpg'))
    idx += 1
    if idx % 100 == 0:
        print idx, 'images done!'

    height, width, channel = img.shape

    neg_num = 0
    while neg_num < 50:
        size = npr.randint(12, min(width, height) / 2)  # 12 or 40
        nx = npr.randint(0, width - size)
        ny = npr.randint(0, height - size)
        crop_box = np.array([nx, ny, nx + size, ny + size])

        iou = IOU(crop_box, bboxes)

        cropped_img = img[ny:ny + size, nx:nx + size, :]
        resized_img = cv2.resize(cropped_img, (12, 12),
                                 interpolation=cv2.INTER_LINEAR)

        if np.max(iou) < 0.3:
            save_file = os.path.join(neg_save_dir, '%s.jpg' % n_idx)
            f2.write('12/negative/%s' % n_idx + ' 0\n')
            cv2.imwrite(save_file, resized_img)
            n_idx += 1
            neg_num += 1
    for box in bboxes:
        # for each ground truth bbox
        x1, y1, x2, y2 = box
        w = x2 - x1 + 1
コード例 #14
0
    scores = []
    tracker = Re3Tracker()
    for video in VIDEO_PATH:
        tracker.reset()
        
        print('Now in folder: ' + video)
        image_path = [video + 'color/' + f for f in os.listdir(video + 'color/')]
        image_path.sort()
        ground_truth = np.loadtxt(video + 'groundtruth.txt', delimiter=',')
        
        img = cv2.imread(image_path[0])
        bbox = to_xyxy(ground_truth[0,:])
        bbox = tracker.track(1, img, bbox)
        
        for i in range(1, len(image_path)):
            img = cv2.imread(image_path[i])
            predicted_bbox = tracker.track(1,img)
            score = IOU.IOU(predicted_bbox, to_xyxy(ground_truth[i,:]))
            print('score:', score)
            scores.append(score)
    
    print('Save scores to file ...')
    
    with open('IOU_scores.txt', 'w') as f:
        for s in scores:
            f.write("%s\n" % s)
    
    print('File saved.')

    
コード例 #15
0
ファイル: simulator.py プロジェクト: yueshen95/re3-pytorch
def render_patch(bbox,
                 background,
                 trackedObjects,
                 cropSize=CROP_SIZE,
                 cropPad=CROP_PAD):
    bboxXYWH = xyxy_to_xywh(bbox)
    image = np.zeros((int(cropSize), int(cropSize), 3), dtype=np.uint8)
    fullBBoxXYWH = bboxXYWH.copy()
    fullBBoxXYWH[[2, 3]] *= cropPad
    fullBBox = xywh_to_xyxy(fullBBoxXYWH)
    fullBBoxXYWH = fullBBoxXYWH
    fullBBoxXYWH[[2, 3]] = np.maximum(fullBBoxXYWH[[2, 3]], 1)
    # First do background
    boxPos = np.array([0, 0, background.shape[1], background.shape[0]])
    boxPosXYWH = xyxy_to_xywh(boxPos)
    cropCoords = np.clip(boxPos - fullBBox[[0, 1, 0, 1]], 0,
                         fullBBoxXYWH[[2, 3, 2, 3]])
    cropCoords *= (cropSize) * 1.0 / fullBBoxXYWH[[2, 3, 2, 3]]
    cropCoords = np.clip(np.round(cropCoords), 0, cropSize).astype(int)

    textureCrop = np.zeros(4)
    textureCrop[0] = int(
        max(fullBBox[0] - boxPos[0], 0) * background.shape[1] * 1.0 /
        boxPosXYWH[2])
    textureCrop[1] = int(
        max(fullBBox[1] - boxPos[1], 0) * background.shape[0] * 1.0 /
        boxPosXYWH[3])
    textureCrop[2] = int(
        min((fullBBox[2] - boxPos[0]) * 1.0 / boxPosXYWH[2], 1) *
        background.shape[1])
    textureCrop[3] = int(
        min((fullBBox[3] - boxPos[1]) * 1.0 / boxPosXYWH[3], 1) *
        background.shape[0])
    if (textureCrop[2] - textureCrop[0] < 1
            or textureCrop[3] - textureCrop[1] < 1):
        textureCrop = [0, 0, 1, 1]
    textureCrop = np.round(textureCrop).astype(int)
    textureCrop[[0, 2]] = np.clip(textureCrop[[0, 2]], 0, background.shape[1])
    textureCrop[[1, 3]] = np.clip(textureCrop[[1, 3]], 0, background.shape[0])

    if cropCoords[3] > cropCoords[1] + 1 and cropCoords[2] > cropCoords[0] + 1:
        image[cropCoords[1]:cropCoords[3],
              cropCoords[0]:cropCoords[2], :] = (cv2.resize(
                  background[textureCrop[1]:textureCrop[3],
                             textureCrop[0]:textureCrop[2], :],
                  (cropCoords[2] - cropCoords[0],
                   cropCoords[3] - cropCoords[1])))

    # Now do all objects
    for obj in trackedObjects:
        boxPos = obj.get_bbox()
        boxPosXYWH = xyxy_to_xywh(boxPos)
        if IOU(boxPos, fullBBox) < 0.001:
            continue
        cropCoords = np.zeros(4)
        cropCoords = np.clip(boxPos - fullBBox[[0, 1, 0, 1]], 0,
                             fullBBoxXYWH[[2, 3, 2, 3]])
        cropCoords *= cropSize * 1.0 / fullBBoxXYWH[[2, 3, 2, 3]]
        cropCoords = np.clip(np.round(cropCoords), 0, cropSize).astype(int)
        if (cropCoords[2] - cropCoords[0] < 1
                or cropCoords[3] - cropCoords[1] < 1):
            cropCoords[[0, 1]] = np.clip(cropCoords[[0, 1]] - 1, 0,
                                         cropSize).astype(int)
            cropCoords[[2, 3]] = np.clip(cropCoords[[2, 3]] + 1, 0,
                                         cropSize).astype(int)
        textureCrop = np.zeros(4, dtype=int)
        textureCrop[0] = int(
            max(fullBBox[0] - boxPos[0], 0) * obj.texture.shape[1] * 1.0 /
            boxPosXYWH[2])
        textureCrop[1] = int(
            max(fullBBox[1] - boxPos[1], 0) * obj.texture.shape[0] * 1.0 /
            boxPosXYWH[3])
        textureCrop[2] = int(
            min((fullBBox[2] - boxPos[0]) * 1.0 / boxPosXYWH[2], 1) *
            obj.texture.shape[1])
        textureCrop[3] = int(
            min((fullBBox[3] - boxPos[1]) * 1.0 / boxPosXYWH[3], 1) *
            obj.texture.shape[0])
        if (textureCrop[2] - textureCrop[0] < 1
                or textureCrop[3] - textureCrop[1] < 1):
            textureCrop = [0, 0, 2, 2]
        textureCrop = np.round(textureCrop).astype(int)
        textureCrop[[0, 2]] = np.clip(textureCrop[[0, 2]], 0,
                                      obj.texture.shape[1])
        textureCrop[[1, 3]] = np.clip(textureCrop[[1, 3]], 0,
                                      obj.texture.shape[0])

        # Feathering
        currentIm = image[cropCoords[1]:cropCoords[3],
                          cropCoords[0]:cropCoords[2], :].astype(np.float32)
        newIm = cv2.resize(
            obj.texture[textureCrop[1]:textureCrop[3],
                        textureCrop[0]:textureCrop[2], :],
            (cropCoords[2] - cropCoords[0], cropCoords[3] - cropCoords[1]),
        ).astype(np.float32)

        if (cropCoords[2] - cropCoords[0] < 1
                or cropCoords[3] - cropCoords[1] < 1):
            featherWeightOn = 0
        else:
            featherCrop = np.zeros(4)
            featherCrop[0] = int(
                max(fullBBox[0] - boxPos[0], 0) *
                FEATHER_WEIGHT_ARRAY.shape[1] * 1.0 / boxPosXYWH[2])
            featherCrop[1] = int(
                max(fullBBox[1] - boxPos[1], 0) *
                FEATHER_WEIGHT_ARRAY.shape[0] * 1.0 / boxPosXYWH[3])
            featherCrop[2] = int(
                min((fullBBox[2] - boxPos[0]) * 1.0 / boxPosXYWH[2], 1) *
                FEATHER_WEIGHT_ARRAY.shape[1])
            featherCrop[3] = int(
                min((fullBBox[3] - boxPos[1]) * 1.0 / boxPosXYWH[3], 1) *
                FEATHER_WEIGHT_ARRAY.shape[0])
            if (featherCrop[2] - featherCrop[0] < 1
                    or featherCrop[3] - featherCrop[1] < 1):
                featherCrop = [
                    int(CROP_SIZE / 2 - 1),
                    int(CROP_SIZE / 2 - 1),
                    int(CROP_SIZE / 2),
                    int(CROP_SIZE / 2)
                ]
            featherCrop = np.round(featherCrop).astype(int)
            featherCrop[[0, 2]] = np.clip(featherCrop[[0, 2]], 0,
                                          FEATHER_WEIGHT_ARRAY.shape[1])
            featherCrop[[1, 3]] = np.clip(featherCrop[[1, 3]], 0,
                                          FEATHER_WEIGHT_ARRAY.shape[0])
            featherWeightOn = cv2.resize(
                FEATHER_WEIGHT_ARRAY[featherCrop[1]:featherCrop[3],
                                     featherCrop[0]:featherCrop[2], :],
                (cropCoords[2] - cropCoords[0],
                 cropCoords[3] - cropCoords[1])).astype(np.float32) / 255.0
        image[cropCoords[1]:cropCoords[3], cropCoords[0]:cropCoords[2], :] = (
            (newIm * featherWeightOn + currentIm *
             (1 - featherWeightOn)).astype(np.uint8))
    return image