def __getitem__(self, idx): fn = self.image_ids[idx] + '.png' img = cv2.imread(path.join(train_dir, fn), cv2.IMREAD_COLOR) img2 = cv2.imread(path.join(train_dir2, fn), cv2.IMREAD_COLOR) img3 = cv2.imread(path.join(train_dir3, fn), cv2.IMREAD_COLOR) msk = cv2.imread(path.join(masks_folder, fn), cv2.IMREAD_COLOR) msk = (msk > 127) * 1 msk = msk[..., :2] img = np.concatenate([img, img2, img3], axis=2) img = img[98:-98, 98:-98, ...] msk = msk[98:-98, 98:-98, ...] img = preprocess_inputs(img) nadir, cat_inp, coord_inp = parse_img_id(fn) img = torch.from_numpy(img.transpose((2, 0, 1))).float() msk = torch.from_numpy(msk.transpose((2, 0, 1)).copy()).long() nadir = torch.from_numpy(np.asarray([nadir / 60.0]).copy()).float() cat_inp = torch.from_numpy(cat_inp.copy()).float() coord_inp = torch.from_numpy(coord_inp.copy()).float() sample = { "img": img, "mask": msk, 'nadir': nadir, 'cat_inp': cat_inp, 'coord_inp': coord_inp, 'img_name': fn } return sample
def __getitem__(self, idx): fn = self.image_ids[idx] + '.png' img = cv2.imread(path.join(train_dir, fn), cv2.IMREAD_COLOR) img2 = cv2.imread(path.join(train_dir2, fn), cv2.IMREAD_COLOR) if img2 is None: print('Error!', fn) img3 = cv2.imread(path.join(train_dir3, fn), cv2.IMREAD_COLOR) msk = cv2.imread(path.join(masks_folder, fn), cv2.IMREAD_COLOR) occluded_msk = cv2.imread(path.join(occluded_masks_dir, fn), cv2.IMREAD_UNCHANGED) if random.random() > 0.6: shift_pnt = (random.randint(-400, 400), random.randint(-400, 400)) img = shift_image(img, shift_pnt) img2 = shift_image(img2, shift_pnt) img3 = shift_image(img3, shift_pnt) msk = shift_image(msk, shift_pnt) occluded_msk = shift_image(occluded_msk, shift_pnt) if random.random() > 0.96: rot_pnt = (img.shape[0] // 2 + random.randint(-150, 150), img.shape[1] // 2 + random.randint(-150, 150)) scale = 1 angle = random.randint(0, 8) - 4 if (angle != 0) or (scale != 1): img = rotate_image(img, angle, scale, rot_pnt) img2 = rotate_image(img2, angle, scale, rot_pnt) img3 = rotate_image(img3, angle, scale, rot_pnt) msk = rotate_image(msk, angle, scale, rot_pnt) occluded_msk = rotate_image(occluded_msk, angle, scale, rot_pnt) crop_size = input_shape[0] if random.random() > 0.8: crop_size = random.randint(int(input_shape[0] / 1.2), int(input_shape[0] / 0.8)) x0 = random.randint(0, img.shape[1] - crop_size) y0 = random.randint(0, img.shape[0] - crop_size) img = img[y0:y0 + crop_size, x0:x0 + crop_size, :] img2 = img2[y0:y0 + crop_size, x0:x0 + crop_size, :] img3 = img3[y0:y0 + crop_size, x0:x0 + crop_size, :] msk = msk[y0:y0 + crop_size, x0:x0 + crop_size, :] occluded_msk = occluded_msk[y0:y0 + crop_size, x0:x0 + crop_size, ...] if crop_size != input_shape[0]: img = cv2.resize(img, input_shape, interpolation=cv2.INTER_LINEAR) img2 = cv2.resize(img2, input_shape, interpolation=cv2.INTER_LINEAR) img3 = cv2.resize(img3, input_shape, interpolation=cv2.INTER_LINEAR) msk = cv2.resize(msk, input_shape, interpolation=cv2.INTER_LINEAR) occluded_msk = cv2.resize(occluded_msk, input_shape, interpolation=cv2.INTER_LINEAR) if random.random() > 0.5: if random.random() > 0.91: img = clahe(img) img2 = clahe(img2) img3 = clahe(img3) elif random.random() > 0.91: img = gauss_noise(img) img2 = gauss_noise(img2) img3 = gauss_noise(img3) elif random.random() > 0.91: img = cv2.blur(img, (3, 3)) img2 = cv2.blur(img2, (3, 3)) img3 = cv2.blur(img3, (3, 3)) else: if random.random() > 0.91: img = saturation(img, 0.9 + random.random() * 0.2) img2 = saturation(img2, 0.9 + random.random() * 0.2) img3 = saturation(img3, 0.9 + random.random() * 0.2) elif random.random() > 0.91: img = brightness(img, 0.9 + random.random() * 0.2) img2 = brightness(img2, 0.9 + random.random() * 0.2) img3 = brightness(img3, 0.9 + random.random() * 0.2) elif random.random() > 0.91: img = contrast(img, 0.9 + random.random() * 0.2) img2 = contrast(img2, 0.9 + random.random() * 0.2) img3 = contrast(img3, 0.9 + random.random() * 0.2) if random.random() > 0.96: el_det = self.elastic.to_deterministic() img = el_det.augment_image(img) img2 = el_det.augment_image(img2) img3 = el_det.augment_image(img3) msk = (msk > 127) * 1 occluded_msk = (occluded_msk > 127) * 1 occluded_msk = occluded_msk[..., np.newaxis] msk = np.concatenate([msk, occluded_msk], axis=2) img = np.concatenate([img, img2, img3], axis=2) img = preprocess_inputs(img) nadir, cat_inp, coord_inp = parse_img_id(fn) img = torch.from_numpy(img.transpose((2, 0, 1))).float() msk = torch.from_numpy(msk.transpose((2, 0, 1)).copy()).long() nadir = torch.from_numpy(np.asarray([nadir / 60.0]).copy()).float() cat_inp = torch.from_numpy(cat_inp.copy()).float() coord_inp = torch.from_numpy(coord_inp.copy()).float() sample = { "img": img, "mask": msk, 'nadir': nadir, 'cat_inp': cat_inp, 'coord_inp': coord_inp, 'img_name': fn } return sample
img3 = cv2.imread(path.join(test_dir3, f), cv2.IMREAD_COLOR) img = np.concatenate([img, img2, img3], axis=2) img = cv2.copyMakeBorder(img, 14, 14, 14, 14, cv2.BORDER_REFLECT_101) inp = [] inp.append(img) inp = np.asarray(inp, dtype='float') inp = preprocess_inputs(inp) inp = torch.from_numpy(inp.transpose((0, 3, 1, 2))).float() inp = Variable(inp).cuda() nadir, cat_inp, coord_inp = parse_img_id(f) nadir = torch.from_numpy(np.asarray([nadir / 60.0]).copy()).float() cat_inp = torch.from_numpy(cat_inp.copy()[np.newaxis, ...]).float() coord_inp = torch.from_numpy(coord_inp.copy()[np.newaxis, ...]).float() msk, _ = model(inp, nadir, cat_inp, coord_inp) msk = torch.sigmoid(msk) msk = msk.cpu().numpy() msk = msk[0] msk = msk * 255 msk = msk.astype('uint8') msk = np.rollaxis(msk, 0, 3)