def find_head_edge(self, box, head):
     head_dict = {0: '11', 1: '10', 2: '00', 3: '01'}
     flag = head_dict[int(head)]
     box[4] += random.random() * 0.1
     box_eight = forward_convert(np.array([box]), False)[0]
     box_eight = np.reshape(box_eight, [4, 2])
     four_edges = [[box_eight[0],
                    box_eight[1]], [box_eight[1], box_eight[2]],
                   [box_eight[2], box_eight[3]],
                   [box_eight[3], box_eight[0]]]
     for i in range(4):
         center_x = (four_edges[i][0][0] + four_edges[i][1][0]) / 2.
         center_y = (four_edges[i][0][1] + four_edges[i][1][1]) / 2.
         if (center_x - box[0]) >= 0 and (center_y - box[1]) >= 0:
             res = '11'
             if res == flag:
                 return four_edges[i]
         elif (center_x - box[0]) >= 0 and (center_y - box[1]) <= 0:
             res = '10'
             if res == flag:
                 return four_edges[i]
         elif (center_x - box[0]) <= 0 and (center_y - box[1]) <= 0:
             res = '00'
             if res == flag:
                 return four_edges[i]
         else:
             res = '01'
             if res == flag:
                 return four_edges[i]
Exemple #2
0
def get_mask(img, boxes):
    boxes = forward_convert(boxes)
    h, w, _ = img.shape
    mask = np.zeros([h, w])
    for b in boxes:
        b = np.reshape(b[0:-1], [4, 2])
        rect = np.array(b, np.int32)
        cv2.fillConvexPoly(mask, rect, 1)
    # mask = cv2.resize(mask, dsize=(h // 16, w // 16))
    mask = np.expand_dims(mask, axis=-1)
    return np.array(mask, np.float32)
    def worker(self, gpu_id, images, det_net, result_queue):
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

        img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
                                                         3])  # is RGB. not BGR
        img_batch = tf.cast(img_plac, tf.float32)

        if self.cfgs.NET_NAME in [
                'resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d'
        ]:
            img_batch = (img_batch / 255 - tf.constant(
                self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD)
        else:
            img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN)

        img_batch = tf.expand_dims(img_batch, axis=0)

        detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
            input_img_batch=img_batch)

        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())

        restorer, restore_ckpt = det_net.get_restorer()

        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

        with tf.Session(config=config) as sess:
            sess.run(init_op)
            if not restorer is None:
                restorer.restore(sess, restore_ckpt)
                print('restore model %d ...' % gpu_id)

            for img_path in images:

                # if 'P0006' not in img_path:
                #     continue

                img = cv2.imread(img_path)

                box_res_rotate = []
                label_res_rotate = []
                score_res_rotate = []

                imgH = img.shape[0]
                imgW = img.shape[1]

                img_short_side_len_list = self.cfgs.IMG_SHORT_SIDE_LEN if isinstance(
                    self.cfgs.IMG_SHORT_SIDE_LEN,
                    list) else [self.cfgs.IMG_SHORT_SIDE_LEN]
                img_short_side_len_list = [
                    img_short_side_len_list[0]
                ] if not self.args.multi_scale else img_short_side_len_list

                if imgH < self.args.h_len:
                    temp = np.zeros([self.args.h_len, imgW, 3], np.float32)
                    temp[0:imgH, :, :] = img
                    img = temp
                    imgH = self.args.h_len

                if imgW < self.args.w_len:
                    temp = np.zeros([imgH, self.args.w_len, 3], np.float32)
                    temp[:, 0:imgW, :] = img
                    img = temp
                    imgW = self.args.w_len

                for hh in range(0, imgH,
                                self.args.h_len - self.args.h_overlap):
                    if imgH - hh - 1 < self.args.h_len:
                        hh_ = imgH - self.args.h_len
                    else:
                        hh_ = hh
                    for ww in range(0, imgW,
                                    self.args.w_len - self.args.w_overlap):
                        if imgW - ww - 1 < self.args.w_len:
                            ww_ = imgW - self.args.w_len
                        else:
                            ww_ = ww
                        src_img = img[hh_:(hh_ + self.args.h_len),
                                      ww_:(ww_ + self.args.w_len), :]

                        for short_size in img_short_side_len_list:
                            max_len = self.cfgs.IMG_MAX_LENGTH
                            if self.args.h_len < self.args.w_len:
                                new_h, new_w = short_size, min(
                                    int(short_size * float(self.args.w_len) /
                                        self.args.h_len), max_len)
                            else:
                                new_h, new_w = min(
                                    int(short_size * float(self.args.h_len) /
                                        self.args.w_len), max_len), short_size
                            img_resize = cv2.resize(src_img, (new_w, new_h))

                            resized_img, det_boxes_r_, det_scores_r_, det_category_r_ = \
                                sess.run(
                                    [img_batch, detection_boxes, detection_scores, detection_category],
                                    feed_dict={img_plac: img_resize[:, :, ::-1]}
                                )

                            resized_h, resized_w = resized_img.shape[
                                1], resized_img.shape[2]
                            src_h, src_w = src_img.shape[0], src_img.shape[1]

                            if len(det_boxes_r_) > 0:
                                det_boxes_r_ = forward_convert(
                                    det_boxes_r_, False)
                                det_boxes_r_[:, 0::2] *= (src_w / resized_w)
                                det_boxes_r_[:, 1::2] *= (src_h / resized_h)

                                for ii in range(len(det_boxes_r_)):
                                    box_rotate = det_boxes_r_[ii]
                                    box_rotate[0::2] = box_rotate[0::2] + ww_
                                    box_rotate[1::2] = box_rotate[1::2] + hh_
                                    box_res_rotate.append(box_rotate)
                                    label_res_rotate.append(
                                        det_category_r_[ii])
                                    score_res_rotate.append(det_scores_r_[ii])

                            if self.args.flip_img:
                                det_boxes_r_flip, det_scores_r_flip, det_category_r_flip = \
                                    sess.run(
                                        [detection_boxes, detection_scores, detection_category],
                                        feed_dict={img_plac: cv2.flip(img_resize, flipCode=1)[:, :, ::-1]}
                                    )
                                if len(det_boxes_r_flip) > 0:
                                    det_boxes_r_flip = forward_convert(
                                        det_boxes_r_flip, False)
                                    det_boxes_r_flip[:, 0::2] *= (src_w /
                                                                  resized_w)
                                    det_boxes_r_flip[:, 1::2] *= (src_h /
                                                                  resized_h)

                                    for ii in range(len(det_boxes_r_flip)):
                                        box_rotate = det_boxes_r_flip[ii]
                                        box_rotate[0::2] = (
                                            src_w - box_rotate[0::2]) + ww_
                                        box_rotate[
                                            1::2] = box_rotate[1::2] + hh_
                                        box_res_rotate.append(box_rotate)
                                        label_res_rotate.append(
                                            det_category_r_flip[ii])
                                        score_res_rotate.append(
                                            det_scores_r_flip[ii])

                                det_boxes_r_flip, det_scores_r_flip, det_category_r_flip = \
                                    sess.run(
                                        [detection_boxes, detection_scores, detection_category],
                                        feed_dict={img_plac: cv2.flip(img_resize, flipCode=0)[:, :, ::-1]}
                                    )
                                if len(det_boxes_r_flip) > 0:
                                    det_boxes_r_flip = forward_convert(
                                        det_boxes_r_flip, False)
                                    det_boxes_r_flip[:, 0::2] *= (src_w /
                                                                  resized_w)
                                    det_boxes_r_flip[:, 1::2] *= (src_h /
                                                                  resized_h)

                                    for ii in range(len(det_boxes_r_flip)):
                                        box_rotate = det_boxes_r_flip[ii]
                                        box_rotate[
                                            0::2] = box_rotate[0::2] + ww_
                                        box_rotate[1::2] = (
                                            src_h - box_rotate[1::2]) + hh_
                                        box_res_rotate.append(box_rotate)
                                        label_res_rotate.append(
                                            det_category_r_flip[ii])
                                        score_res_rotate.append(
                                            det_scores_r_flip[ii])

                box_res_rotate = np.array(box_res_rotate)
                label_res_rotate = np.array(label_res_rotate)
                score_res_rotate = np.array(score_res_rotate)

                box_res_rotate_ = []
                label_res_rotate_ = []
                score_res_rotate_ = []
                threshold = {
                    'roundabout': 0.1,
                    'tennis-court': 0.3,
                    'swimming-pool': 0.1,
                    'storage-tank': 0.2,
                    'soccer-ball-field': 0.3,
                    'small-vehicle': 0.2,
                    'ship': 0.2,
                    'plane': 0.3,
                    'large-vehicle': 0.1,
                    'helicopter': 0.2,
                    'harbor': 0.0001,
                    'ground-track-field': 0.3,
                    'bridge': 0.0001,
                    'basketball-court': 0.3,
                    'baseball-diamond': 0.3
                }

                for sub_class in range(1, self.cfgs.CLASS_NUM + 1):
                    index = np.where(label_res_rotate == sub_class)[0]
                    if len(index) == 0:
                        continue
                    tmp_boxes_r = box_res_rotate[index]
                    tmp_label_r = label_res_rotate[index]
                    tmp_score_r = score_res_rotate[index]

                    tmp_boxes_r_ = backward_convert(tmp_boxes_r, False)

                    try:
                        inx = nms_rotate.nms_rotate_cpu(
                            boxes=np.array(tmp_boxes_r_),
                            scores=np.array(tmp_score_r),
                            iou_threshold=threshold[
                                self.label_name_map[sub_class]],
                            max_output_size=5000)
                    except:
                        tmp_boxes_r_ = np.array(tmp_boxes_r_)
                        tmp = np.zeros(
                            [tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
                        tmp[:, 0:-1] = tmp_boxes_r_
                        tmp[:, -1] = np.array(tmp_score_r)
                        # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
                        jitter = np.zeros(
                            [tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
                        jitter[:, 0] += np.random.rand(
                            tmp_boxes_r_.shape[0], ) / 1000
                        inx = rotate_gpu_nms(
                            np.array(tmp, np.float32) +
                            np.array(jitter, np.float32),
                            float(threshold[self.label_name_map[sub_class]]),
                            0)

                    box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
                    score_res_rotate_.extend(np.array(tmp_score_r)[inx])
                    label_res_rotate_.extend(np.array(tmp_label_r)[inx])

                result_dict = {
                    'boxes': np.array(box_res_rotate_),
                    'scores': np.array(score_res_rotate_),
                    'labels': np.array(label_res_rotate_),
                    'image_id': img_path
                }
                result_queue.put_nowait(result_dict)
    def eval_with_plac(self, img_dir, det_net, image_ext):

        os.environ["CUDA_VISIBLE_DEVICES"] = self.args.gpu
        # 1. preprocess img
        img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
                                                         3])  # is RGB. not BGR
        img_batch = tf.cast(img_plac, tf.float32)

        pretrain_zoo = PretrainModelZoo()
        if self.cfgs.NET_NAME in pretrain_zoo.pth_zoo or self.cfgs.NET_NAME in pretrain_zoo.mxnet_zoo:
            img_batch = (img_batch / 255 - tf.constant(
                self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD)
        else:
            img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN)

        img_batch = tf.expand_dims(img_batch, axis=0)

        detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
            input_img_batch=img_batch)

        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())

        restorer, restore_ckpt = det_net.get_restorer()

        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

        with tf.Session(config=config) as sess:
            sess.run(init_op)
            if not restorer is None:
                restorer.restore(sess, restore_ckpt)
                print('restore model')

            all_boxes_r = []
            imgs = os.listdir(img_dir)
            pbar = tqdm(imgs)
            for a_img_name in pbar:
                a_img_name = a_img_name.split(image_ext)[0]

                raw_img = cv2.imread(
                    os.path.join(img_dir, a_img_name + image_ext))
                raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]

                det_boxes_r_all, det_scores_r_all, det_category_r_all = [], [], []

                img_short_side_len_list = self.cfgs.IMG_SHORT_SIDE_LEN if isinstance(
                    self.cfgs.IMG_SHORT_SIDE_LEN,
                    list) else [self.cfgs.IMG_SHORT_SIDE_LEN]
                img_short_side_len_list = [
                    img_short_side_len_list[0]
                ] if not self.args.multi_scale else img_short_side_len_list

                for short_size in img_short_side_len_list:
                    max_len = self.cfgs.IMG_MAX_LENGTH
                    if raw_h < raw_w:
                        new_h, new_w = short_size, min(
                            int(short_size * float(raw_w) / raw_h), max_len)
                    else:
                        new_h, new_w = min(
                            int(short_size * float(raw_h) / raw_w),
                            max_len), short_size
                    img_resize = cv2.resize(raw_img, (new_w, new_h))

                    resized_img, detected_boxes, detected_scores, detected_categories = \
                        sess.run(
                            [img_batch, detection_boxes, detection_scores, detection_category],
                            feed_dict={img_plac: img_resize[:, :, ::-1]}
                        )

                    if detected_boxes.shape[0] == 0:
                        continue
                    resized_h, resized_w = resized_img.shape[
                        1], resized_img.shape[2]
                    detected_boxes = forward_convert(detected_boxes, False)
                    detected_boxes[:, 0::2] *= (raw_w / resized_w)
                    detected_boxes[:, 1::2] *= (raw_h / resized_h)

                    det_boxes_r_all.extend(detected_boxes)
                    det_scores_r_all.extend(detected_scores)
                    det_category_r_all.extend(detected_categories)
                det_boxes_r_all = np.array(det_boxes_r_all)
                det_scores_r_all = np.array(det_scores_r_all)
                det_category_r_all = np.array(det_category_r_all)

                box_res_rotate_ = []
                label_res_rotate_ = []
                score_res_rotate_ = []

                if det_scores_r_all.shape[0] != 0:
                    for sub_class in range(1, self.cfgs.CLASS_NUM + 1):
                        index = np.where(det_category_r_all == sub_class)[0]
                        if len(index) == 0:
                            continue
                        tmp_boxes_r = det_boxes_r_all[index]
                        tmp_label_r = det_category_r_all[index]
                        tmp_score_r = det_scores_r_all[index]

                        if self.args.multi_scale:
                            tmp_boxes_r_ = backward_convert(tmp_boxes_r, False)

                            # try:
                            #     inx = nms_rotate.nms_rotate_cpu(boxes=np.array(tmp_boxes_r_),
                            #                                     scores=np.array(tmp_score_r),
                            #                                     iou_threshold=self.cfgs.NMS_IOU_THRESHOLD,
                            #                                     max_output_size=5000)
                            # except:
                            tmp_boxes_r_ = np.array(tmp_boxes_r_)
                            tmp = np.zeros([
                                tmp_boxes_r_.shape[0],
                                tmp_boxes_r_.shape[1] + 1
                            ])
                            tmp[:, 0:-1] = tmp_boxes_r_
                            tmp[:, -1] = np.array(tmp_score_r)
                            # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
                            jitter = np.zeros([
                                tmp_boxes_r_.shape[0],
                                tmp_boxes_r_.shape[1] + 1
                            ])
                            jitter[:, 0] += np.random.rand(
                                tmp_boxes_r_.shape[0], ) / 1000
                            inx = rotate_gpu_nms(
                                np.array(tmp, np.float32) +
                                np.array(jitter, np.float32),
                                float(self.cfgs.NMS_IOU_THRESHOLD), 0)
                        else:
                            inx = np.arange(0, tmp_score_r.shape[0])

                        box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
                        score_res_rotate_.extend(np.array(tmp_score_r)[inx])
                        label_res_rotate_.extend(np.array(tmp_label_r)[inx])

                if len(box_res_rotate_) == 0:
                    all_boxes_r.append(np.array([]))
                    continue

                det_boxes_r_ = np.array(box_res_rotate_)
                det_scores_r_ = np.array(score_res_rotate_)
                det_category_r_ = np.array(label_res_rotate_)

                if self.args.draw_imgs:
                    detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE
                    detected_scores = det_scores_r_[detected_indices]
                    detected_boxes = det_boxes_r_[detected_indices]
                    detected_categories = det_category_r_[detected_indices]

                    detected_boxes = backward_convert(detected_boxes, False)

                    drawer = DrawBox(self.cfgs)

                    det_detections_r = drawer.draw_boxes_with_label_and_scores(
                        raw_img[:, :, ::-1],
                        boxes=detected_boxes,
                        labels=detected_categories,
                        scores=detected_scores,
                        method=1,
                        in_graph=True)

                    save_dir = os.path.join('test_hrsc', self.cfgs.VERSION,
                                            'hrsc2016_img_vis')
                    tools.makedirs(save_dir)

                    cv2.imwrite(save_dir + '/{}.jpg'.format(a_img_name),
                                det_detections_r[:, :, ::-1])

                det_boxes_r_ = backward_convert(det_boxes_r_, False)

                x_c, y_c, w, h, theta = det_boxes_r_[:, 0], det_boxes_r_[:, 1], det_boxes_r_[:, 2], \
                                        det_boxes_r_[:, 3], det_boxes_r_[:, 4]

                boxes_r = np.transpose(np.stack([x_c, y_c, w, h, theta]))
                dets_r = np.hstack((det_category_r_.reshape(-1, 1),
                                    det_scores_r_.reshape(-1, 1), boxes_r))
                all_boxes_r.append(dets_r)

                pbar.set_description("Eval image %s" % a_img_name)

            # fw1 = open(cfgs.VERSION + '_detections_r.pkl', 'wb')
            # pickle.dump(all_boxes_r, fw1)
            return all_boxes_r
    def worker(self, gpu_id, images, det_net, result_queue):
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
        # 1. preprocess img
        img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3])  # is RGB. not BGR
        img_batch = tf.cast(img_plac, tf.float32)

        if self.cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d',
                                  'resnet152_v1b', 'resnet101_v1b', 'resnet50_v1b', 'resnet34_v1b', 'resnet18_v1b']:
            img_batch = (img_batch / 255 - tf.constant(self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD)
        else:
            img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN)

        img_batch = tf.expand_dims(img_batch, axis=0)

        detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
            input_img_batch=img_batch,
            gtboxes_batch_h=None,
            gtboxes_batch_r=None,
            gpu_id=0)

        init_op = tf.group(
            tf.global_variables_initializer(),
            tf.local_variables_initializer()
        )

        restorer, restore_ckpt = det_net.get_restorer()

        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

        with tf.Session(config=config) as sess:
            sess.run(init_op)
            if not restorer is None:
                restorer.restore(sess, restore_ckpt)
                print('restore model %d ...' % gpu_id)
            for a_img in images:
                raw_img = cv2.imread(a_img)
                raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]

                det_boxes_r_all, det_scores_r_all, det_category_r_all = [], [], []

                img_short_side_len_list = self.cfgs.IMG_SHORT_SIDE_LEN if isinstance(self.cfgs.IMG_SHORT_SIDE_LEN, list) else [
                    self.cfgs.IMG_SHORT_SIDE_LEN]
                img_short_side_len_list = [img_short_side_len_list[0]] if not self.args.multi_scale else img_short_side_len_list

                for short_size in img_short_side_len_list:
                    max_len = self.cfgs.IMG_MAX_LENGTH
                    if raw_h < raw_w:
                        new_h, new_w = short_size, min(int(short_size * float(raw_w) / raw_h), max_len)
                    else:
                        new_h, new_w = min(int(short_size * float(raw_h) / raw_w), max_len), short_size
                    img_resize = cv2.resize(raw_img, (new_w, new_h))

                    resized_img, detected_boxes, detected_scores, detected_categories = \
                        sess.run(
                            [img_batch, detection_boxes, detection_scores, detection_category],
                            feed_dict={img_plac: img_resize[:, :, ::-1]}
                        )

                    detected_indices = detected_scores >= self.cfgs.VIS_SCORE
                    detected_scores = detected_scores[detected_indices]
                    detected_boxes = detected_boxes[detected_indices]
                    detected_categories = detected_categories[detected_indices]

                    if detected_boxes.shape[0] == 0:
                        continue
                    resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
                    detected_boxes = forward_convert(detected_boxes, False)
                    detected_boxes[:, 0::2] *= (raw_w / resized_w)
                    detected_boxes[:, 1::2] *= (raw_h / resized_h)

                    det_boxes_r_all.extend(detected_boxes)
                    det_scores_r_all.extend(detected_scores)
                    det_category_r_all.extend(detected_categories)

                    if self.args.flip_img:
                        detected_boxes, detected_scores, detected_categories = \
                            sess.run(
                                [detection_boxes, detection_scores, detection_category],
                                feed_dict={img_plac: cv2.flip(img_resize, flipCode=1)[:, :, ::-1]}
                            )
                        detected_indices = detected_scores >= self.cfgs.VIS_SCORE
                        detected_scores = detected_scores[detected_indices]
                        detected_boxes = detected_boxes[detected_indices]
                        detected_categories = detected_categories[detected_indices]

                        if detected_boxes.shape[0] == 0:
                            continue
                        resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
                        detected_boxes = forward_convert(detected_boxes, False)
                        detected_boxes[:, 0::2] *= (raw_w / resized_w)
                        detected_boxes[:, 0::2] = (raw_w - detected_boxes[:, 0::2])
                        detected_boxes[:, 1::2] *= (raw_h / resized_h)

                        det_boxes_r_all.extend(sort_corners(detected_boxes))
                        det_scores_r_all.extend(detected_scores)
                        det_category_r_all.extend(detected_categories)

                        detected_boxes, detected_scores, detected_categories = \
                            sess.run(
                                [detection_boxes, detection_scores, detection_category],
                                feed_dict={img_plac: cv2.flip(img_resize, flipCode=0)[:, :, ::-1]}
                            )
                        detected_indices = detected_scores >= self.cfgs.VIS_SCORE
                        detected_scores = detected_scores[detected_indices]
                        detected_boxes = detected_boxes[detected_indices]
                        detected_categories = detected_categories[detected_indices]

                        if detected_boxes.shape[0] == 0:
                            continue
                        resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
                        detected_boxes = forward_convert(detected_boxes, False)
                        detected_boxes[:, 0::2] *= (raw_w / resized_w)
                        detected_boxes[:, 1::2] *= (raw_h / resized_h)
                        detected_boxes[:, 1::2] = (raw_h - detected_boxes[:, 1::2])
                        det_boxes_r_all.extend(sort_corners(detected_boxes))
                        det_scores_r_all.extend(detected_scores)
                        det_category_r_all.extend(detected_categories)

                det_boxes_r_all = np.array(det_boxes_r_all)
                det_scores_r_all = np.array(det_scores_r_all)
                det_category_r_all = np.array(det_category_r_all)

                box_res_rotate_ = []
                label_res_rotate_ = []
                score_res_rotate_ = []

                if det_scores_r_all.shape[0] != 0:
                    for sub_class in range(1, self.cfgs.CLASS_NUM + 1):
                        index = np.where(det_category_r_all == sub_class)[0]
                        if len(index) == 0:
                            continue
                        tmp_boxes_r = det_boxes_r_all[index]
                        tmp_label_r = det_category_r_all[index]
                        tmp_score_r = det_scores_r_all[index]

                        if self.args.multi_scale:
                            tmp_boxes_r_ = backward_convert(tmp_boxes_r, False)

                            # try:
                            #     inx = nms_rotate.nms_rotate_cpu(boxes=np.array(tmp_boxes_r_),
                            #                                     scores=np.array(tmp_score_r),
                            #                                     iou_threshold=self.cfgs.NMS_IOU_THRESHOLD,
                            #                                     max_output_size=5000)
                            # except:
                            tmp_boxes_r_ = np.array(tmp_boxes_r_)
                            tmp = np.zeros([tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
                            tmp[:, 0:-1] = tmp_boxes_r_
                            tmp[:, -1] = np.array(tmp_score_r)
                            # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
                            jitter = np.zeros([tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
                            jitter[:, 0] += np.random.rand(tmp_boxes_r_.shape[0], ) / 1000
                            inx = rotate_gpu_nms(np.array(tmp, np.float32) + np.array(jitter, np.float32),
                                                 float(self.cfgs.NMS_IOU_THRESHOLD), 0)
                        else:
                            inx = np.arange(0, tmp_score_r.shape[0])

                        box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
                        score_res_rotate_.extend(np.array(tmp_score_r)[inx])
                        label_res_rotate_.extend(np.array(tmp_label_r)[inx])

                box_res_rotate_ = np.array(box_res_rotate_)
                score_res_rotate_ = np.array(score_res_rotate_)
                label_res_rotate_ = np.array(label_res_rotate_)

                result_dict = {'scales': [1, 1], 'boxes': box_res_rotate_,
                               'scores': score_res_rotate_, 'labels': label_res_rotate_,
                               'image_id': a_img}
                result_queue.put_nowait(result_dict)
Exemple #6
0
                temp_ious.append(0.0)
            ious.append(temp_ious)

    return np.array(ious, dtype=np.float32)


if __name__ == '__main__':
    boxes1 = np.array([[50, 50, 70, 10, -45]], np.float32)

    boxes2 = np.array([[50, 25, 70, 10, -45]], np.float32)

    iou_r = riou(boxes1, boxes2)
    print(iou_r)

    from libs.utils.coordinate_convert import forward_convert
    boxes1 = forward_convert(boxes1, False)
    boxes2 = forward_convert(boxes2, False)

    x1_min = np.min(boxes1[:, ::2], axis=-1)
    y1_min = np.min(boxes1[:, 1::2], axis=-1)
    x1_max = np.max(boxes1[:, ::2], axis=-1)
    y1_max = np.max(boxes1[:, 1::2], axis=-1)
    boxes1 = np.transpose(np.stack([x1_min, y1_min, x1_max, y1_max]))

    x2_min = np.min(boxes2[:, ::2], axis=-1)
    y2_min = np.min(boxes2[:, 1::2], axis=-1)
    x2_max = np.max(boxes2[:, ::2], axis=-1)
    y2_max = np.max(boxes2[:, 1::2], axis=-1)
    boxes2 = np.transpose(np.stack([x2_min, y2_min, x2_max, y2_max]))

    iou_h = hiou(boxes1, boxes2)