From 8189b6005777649bf7a601eece4cff5ccfc7266e Mon Sep 17 00:00:00 2001 From: jwangzzz <j96w@qq.com> Date: Thu, 19 Sep 2019 13:07:17 -0400 Subject: [PATCH] fix the bbox bug --- datasets/linemod/dataset.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/datasets/linemod/dataset.py b/datasets/linemod/dataset.py index a5aacdc..3f667d6 100755 --- a/datasets/linemod/dataset.py +++ b/datasets/linemod/dataset.py @@ -18,6 +18,7 @@ import copy import scipy.misc import scipy.io as scio import yaml +import cv2 class PoseDataset(data.Dataset): @@ -117,7 +118,10 @@ class PoseDataset(data.Dataset): img = np.transpose(img, (2, 0, 1)) img_masked = img - rmin, rmax, cmin, cmax = get_bbox(meta['obj_bb']) + if self.mode == 'eval': + rmin, rmax, cmin, cmax = get_bbox(mask_to_bbox(mask_label)) + else: + rmin, rmax, cmin, cmax = get_bbox(meta['obj_bb']) img_masked = img_masked[:, rmin:rmax, cmin:cmax] #p_img = np.transpose(img_masked, (1, 2, 0)) @@ -209,6 +213,24 @@ border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520 img_width = 480 img_length = 640 + +def mask_to_bbox(mask): + mask = mask.astype(np.uint8) + _, contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + x = 0 + y = 0 + w = 0 + h = 0 + for contour in contours: + tmp_x, tmp_y, tmp_w, tmp_h = cv2.boundingRect(contour) + if tmp_w * tmp_h > w * h: + x = tmp_x + y = tmp_y + w = tmp_w + h = tmp_h + return [x, y, w, h] + + def get_bbox(bbox): bbx = [bbox[1], bbox[1] + bbox[3], bbox[0], bbox[0] + bbox[2]] if bbx[0] < 0: -- GitLab