| 51 | self._load(name) |
| 52 | |
| 53 | def _load(self, name): |
| 54 | image_glob = os.path.join(self.data_root, 'images', name, '*.jpg') |
| 55 | image_files = glob.glob(image_glob) |
| 56 | gt_dir = os.path.join(self.data_root, 'groundTruth', name) |
| 57 | self.data = np.zeros((len(image_files), IMG_H, IMG_W, 3), dtype='uint8') |
| 58 | self.label = np.zeros((len(image_files), IMG_H, IMG_W), dtype='float32') |
| 59 | |
| 60 | for idx, f in enumerate(image_files): |
| 61 | im = cv2.imread(f, cv2.IMREAD_COLOR) |
| 62 | assert im is not None |
| 63 | if im.shape[0] > im.shape[1]: |
| 64 | im = np.transpose(im, (1, 0, 2)) |
| 65 | assert im.shape[:2] == (IMG_H, IMG_W), "{} != {}".format(im.shape[:2], (IMG_H, IMG_W)) |
| 66 | |
| 67 | imgid = os.path.basename(f).split('.')[0] |
| 68 | gt_file = os.path.join(gt_dir, imgid) |
| 69 | gt = loadmat(gt_file)['groundTruth'][0] |
| 70 | n_annot = gt.shape[0] |
| 71 | gt = sum(gt[k]['Boundaries'][0][0] for k in range(n_annot)) |
| 72 | gt = gt.astype('float32') |
| 73 | gt *= 1.0 / n_annot |
| 74 | if gt.shape[0] > gt.shape[1]: |
| 75 | gt = gt.transpose() |
| 76 | assert gt.shape == (IMG_H, IMG_W) |
| 77 | |
| 78 | self.data[idx] = im |
| 79 | self.label[idx] = gt |
| 80 | |
| 81 | def __len__(self): |
| 82 | return self.data.shape[0] |