def __getitem__(self, idx):
first_idx = idx if idx % config.num_files_folder == 0 else idx - 1
next_idx = idx
file1, file2 = self.filenames[first_idx], self.filenames[next_idx]
first = cv2.cvtColor(cv2.imread(str(file1)), cv2.COLOR_BGR2RGB)
second = cv2.cvtColor(cv2.imread(str(file2)), cv2.COLOR_BGR2RGB)
return str(file2), img_to_tensor(first), img_to_tensor(second)
def run(self):
self.game.init_game()
state = self.game.start_game()
episode_reward = 0
while True:
logger.info("Next frame.")
epsilon, learning = self.learner.epsilon
# all learning iterations done
if not learning:
break
# get best or random action with probability epsilon
action = self.learner.act(state, epsilon)
next_state, reward, done = self.game.move(action)
self.learner.remember(img_to_tensor(state), action, reward,
img_to_tensor(next_state), done)
state = next_state
episode_reward += reward
if done:
logger.warning("Game over ({} fps).".format(
len(self.game.frames) / (time.time() - self.game.started)))
self.learner.update()
self.learner.add_episode_reward(episode_reward)
episode_reward = 0
# reset
self.game.reset()
state = self.game.start_game()
def __getitem__(self, idx):
filename = self.filenames[idx]
image = load_image(filename)
mask = load_mask(filename)
data = {'image': image, 'mask': mask}
augmented = self.transform(**data)
image, mask = augmented["image"], augmented["mask"]
# load optical flow
optflow = load_optflow(filename)
optflow = cv2.resize(optflow,
dsize=(image.shape[1], image.shape[0]),
interpolation=cv2.INTER_AREA).transpose(2, 0, 1)
# in evaluation mode should return the filename
if self.mode == 'eval':
return str(filename), img_to_tensor(image), torch.from_numpy(
optflow).float()
# TODO: change the file format
if params.problem_type == 'binary':
# reshape the binary output to be 3D
return img_to_tensor(image), torch.from_numpy(np.expand_dims(mask, 0)).float(), \
torch.from_numpy(optflow).float()
else:
return img_to_tensor(image), torch.from_numpy(mask).long(), \
torch.from_numpy(optflow).float()
def __getitem__(self, idx):
input_filename = self.input_filenames[idx]
label_filename = self.label_filenames[idx]
input_image = load_image(input_filename)
label_image = load_image(label_filename)
augmented_input_image, augmented_label_image = self.transform(
image=input_image, mask=label_image)
return img_to_tensor(augmented_input_image), img_to_tensor(
augmented_label_image)
def __getitem__(self, idx):
first_idx = idx if idx == 0 else idx - 1
next_idx = idx
file1, file2 = self.filenames[first_idx], self.filenames[next_idx]
first = cv2.cvtColor(cv2.imread(str(file1)), cv2.COLOR_BGR2RGB)
second = cv2.cvtColor(cv2.imread(str(file2)), cv2.COLOR_BGR2RGB)
# first = cv2.resize(first, dsize=(1280, 384))
# second = cv2.resize(second, dsize=(1280, 384))
return str(file2), img_to_tensor(first), img_to_tensor(second)
def __getitem__(self, idx):
rec_left_path = self.rec_imgs_and_gts_names['Rectified_left_image_path'].values[idx]
rec_right_path = self.rec_imgs_and_gts_names['Rectified_right_image_path'].values[idx]
rec_left_gt_path = self.rec_imgs_and_gts_names['Rectified_left_gt_path'].values[idx]
# rec_right_gt_path = self.rec_imgs_and_gts_names['Rectified_right_gt_path'].values[idx]
# mask_left_path = self.rec_imgs_and_gts_names['Rectified_left_mask_path'].values[idx]
# mask_right_path = self.rec_imgs_and_gts_names['Rectified_right_mask_path'].values[idx]
reprojection_matrix_left = self.reprojecton_matrixs['arr_0'][idx]
# reprojection_matrix_right = reprojection_matrix_left.copy()
# reprojection_matrix_right[3][2] = - reprojection_matrix_right[3][2]
# bb = self.baseline[idx]
# ff = self.focal[idx]
rec_left = cv2.imread(rec_left_path)
rec_right = cv2.imread(rec_right_path)
rec_left_gt = tifffile.imread(rec_left_gt_path)
# rec_right_gt = tifffile.imread(rec_right_gt_path)
ds_left = cv2.resize(rec_left, (0, 0), fx=1. / self.downsampling, fy=1. / self.downsampling)
ds_right = cv2.resize(rec_right, (0, 0), fx=1. / self.downsampling, fy=1. / self.downsampling)
del rec_left
del rec_right
# ds_left = ds_left.astype(np.float32)
# ds_right = ds_right.astype(np.float32)
rec_left_gt = rec_left_gt.astype(np.float32)
# rec_right_gt = rec_right_gt.astype(np.float32)
left_z = rec_left_gt[:, :, 2]
# right_z = rec_right_gt[:, :, 2]
med_left = np.median(left_z[left_z > 0])
# med_right = np.median(right_z[right_z > 0])
mask_left = (left_z > 5).astype(np.float32) * (left_z < 5 * med_left).astype(np.float32)
# mask_right = (right_z > 5).astype(np.float32) * (right_z < 5 * med_right).astype(np.float32)
mask_left = mask_left[..., np.newaxis]
# mask_right = mask_right[..., np.newaxis]
mask_left = np.repeat(mask_left, 3, axis=2)
# mask_right = np.repeat(mask_right, 3, axis=2)
# bb和ff转换成array，类型为float32
reprojection_matrix_left = reprojection_matrix_left.astype(np.float32)
# reprojection_matrix_right = reprojection_matrix_right.astype(np.float32)
# bb = np.array(bb).astype(np.float32)
# ff = np.array(ff).astype(np.float32)
if self.phase == 'train':
if self.transform is not None:
ds_left = self.transform(image=ds_left)['image']
ds_right = self.transform(image=ds_right)['image']
ds_left = self.normalize(image=ds_left)['image']
ds_right = self.normalize(image=ds_right)['image']
else:
ds_left = self.normalize(image=ds_left)['image']
ds_right = self.normalize(image=ds_right)['image']
return [img_to_tensor(ds_left), img_to_tensor(ds_right), img_to_tensor(rec_left_gt),\
img_to_tensor(mask_left), torch.from_numpy(reprojection_matrix_left)]
def __getitem__(self, idx):
# an exception for each video:
# for the first frame in each video, the optical flow should be 0
# optflow[0] = flow<0, 0> = 0, optflow[k] = flow<k-1, k> (k > 0)
first_idx = idx if idx % utils.num_frames_video == 0 else idx - 1
next_idx = idx
file1, file2 = self.filenames[first_idx], self.filenames[next_idx]
# according to the UnFlow implementation, inputs are in normalized BGR space
first = cv2.imread(str(file1))
second = cv2.imread(str(file2))
# img_to_tensor will reshape into (c, h, w) and scaled to [0., 1.]
return str(file2), img_to_tensor(first), img_to_tensor(second)
def __getitem__(self, idx):
# select the selected file
if self.schedule == "shuffle":
filename = self.shuffled_filenames[idx]
idx = self.shuffled_idx[idx]
elif self.schedule == "ordered":
filename = self.ordered_filenames[idx]
idx = self.ordered_idx[idx]
else:
filename = self.filenames[idx]
self.is_labeled = self.labeled[idx]
# load optical flow
optflow = self.load_optflow(filename)
# print(filename)
image = self.load_image(filename)
if self.is_labeled == False and self.mode == 'train':
# give last frame prediction
mask = self.load_mask(self.filenames[idx - 1])
# mask = self.dilated_mask(mask, optflow)
else:
mask = self.load_mask(filename)
data = {'image': image, 'mask': mask}
augmented = self.transform(**data)
image, mask = augmented["image"], augmented["mask"]
# generate attention map
if idx % config.num_files_folder:
# this is not the first frame, calculate attention map using previous prediction
attmap = self.cal_attmap(self.attmaps[idx - 1], optflow)
else:
# else use last epoch prediction
attmap = self.attmaps[idx]
# in evaluation mode should return the filename
if self.mode == 'eval':
return idx, str(filename), img_to_tensor(image), torch.from_numpy(np.expand_dims(attmap, 0)).float() \
, torch.from_numpy(optflow.transpose(2,0,1)).float()
# TODO: change the file format
if self.problem_type == 'binary':
# reshape the binary output to be 3D
return idx, self.is_labeled, img_to_tensor(image), torch.from_numpy(np.expand_dims(mask, 0)).float(), \
torch.from_numpy(np.expand_dims(attmap, 0)).float(), torch.from_numpy(optflow.transpose(2,0,1)).float()
else:
return idx, self.is_labeled, img_to_tensor(image), torch.from_numpy(mask).long(), \
torch.from_numpy(np.expand_dims(attmap, 0)).float(), torch.from_numpy(optflow.transpose(2,0,1)).float()
def __getitem__(self, index):
#image_path = self.image_paths[index]
#image = cv2.imread(image_path, cv2.IMREAD_COLOR)
#image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = self.images[index]
if self.transforms:
image = self.transforms(image=image)['image']
if self.use_masking:
mask_num = random.randrange(1, args.max_mask_num+1)
if self.args.masking_type == 'soft':
mask_prob = 0.25
rand_val = random.random()
if rand_val < mask_prob:
image = time_mask(image, num_masks=mask_num)
rand_val = random.random()
if rand_val < mask_prob:
image = signal_mask(image, num_masks=mask_num)
elif self.args.masking_type == 'hard':
mask_prob = 0.3
rand_val = random.random()
if rand_val < mask_prob:
image = signal_mask(time_mask(image, num_masks=mask_num), num_masks=mask_num)
else:
rand_val = random.random()
if rand_val < 0.3:
image = time_mask(image, num_masks=mask_num)
rand_val = random.random()
if rand_val < 0.3:
image = signal_mask(image, num_masks=mask_num)
#image = img_to_tensor(image, {"mean": [0.485, 0.456, 0.406],
# "std": [0.229, 0.224, 0.225]})
#image = torch.from_numpy(image.transpose((2, 0, 1)))
data = {}
if self.use_meta:
data['image'] = img_to_tensor(image)
data['meta'] = torch.tensor(self.meta[index]).float()
else:
data['image'] = img_to_tensor(image)
if self.is_test:
#return image #torch.tensor(img, dtype=torch.float32)
return data
else:
label = self.labels[index]
#return image, label#torch.tensor(img, dtype=torch.float32),\
#torch.tensor(self.labels[index], dtype=torch.long)
return data, label
def __getitem__(self, idx):
prev_idx = idx
next_idxs = []
seconds = []
for i in range(idx + 1, idx + self.batch_size):
next_idx = i if i < len(
self.filenames) else len(self.filenames) - 1
file1, file2 = self.filenames[prev_idx], self.filenames[next_idx]
first = cv2.cvtColor(cv2.imread(str(file1)), cv2.COLOR_BGR2RGB)
second = cv2.cvtColor(cv2.imread(str(file2)), cv2.COLOR_BGR2RGB)
next_idxs.append(next_idx)
seconds.append(img_to_tensor(second))
return str(file1), prev_idx, img_to_tensor(first), next_idxs, seconds
def __getitem__(self, idx):
img_file_name = self.file_names[idx]
image = load_image(img_file_name)
mask = load_mask(img_file_name, self.problem_type)
data = {"image": image, "mask": mask}
if self.transform is not None:
augmented = self.transform(**data)
image, mask = augmented["image"], augmented["mask"]
if self.problem_type == 'binary':
return img_to_tensor(image), torch.from_numpy(mask).long()
else:
return img_to_tensor(image), torch.from_numpy(mask).long()
def predict(self, image_path):
'''
模型预测返回结果
:param input: 评估传入样例 {"image_path": "./data/input/cloudy/00000.jpg"}
:return: 模型预测成功中户 {"label": 0}
'''
print(image_path)
# Normalize = {'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225)}
result = []
for size in SIZES:
img = cv.imread(image_path)
img = cv.resize(img, dsize=(size, size))
tensor = img_to_tensor(img)
tensor = Variable(torch.unsqueeze(tensor, dim=0).float(),
requires_grad=False)
for subModel in self.models:
output = tta.fliplr_image2label(subModel, tensor.to(device))
result.append(output)
new_output = torch.mean(torch.stack(result, 0), 0)
pred = new_output.max(1, keepdim=True)[1]
# output = tta.fliplr_image2label(self.model, tensor.to(device))
# pred = output.max(1, keepdim=True)[1].item()
return {"label": pred}
def __getitem__(self, idx):
imageid = self.image_ids[idx]
im = get_image(imageid, basepath=self.basepath, rgbdir='test_rgb')
assert im is not None
augmented = self.aug(image=im)
return img_to_tensor(augmented['image']), imageid
def __getitem__(self, index):
image_path = self.image_paths[index]
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
if self.transforms:
image = self.transforms(image=image)['image']
# Masking Augmentation (Saewon)
if self.masking == 'soft':
mask_prob = 0.3
rand_val = random.random()
if rand_val < mask_prob:
image = time_mask(image, num_masks=2)
rand_val = random.random()
if rand_val < mask_prob:
image = signal_mask(image, num_masks=2)
elif self.masking == 'hard':
mask_prob = 0.5
rand_val = random.random()
if rand_val < mask_prob:
image = signal_mask(time_mask(image, num_masks=2), num_masks=2)
image = img_to_tensor(image, {
"mean": [0.485, 0.456, 0.406],
"std": [0.229, 0.224, 0.225]
})
if self.is_test:
return image #torch.tensor(img, dtype=torch.float32)
else:
label = self.labels[index]
return image, label #torch.tensor(img, dtype=torch.float32),\
def __getitem__(self, idx):
'''
pytorch dataloader get_item_from_index
input:
idx: corresponding with __len__
output:
input_dict: a dictionary stores all return value
'''
# input dict for return
input_dict = {}
filename, abs_idx = self.get_filename_from_idx(idx)
image = self.load_image(filename)
# extra input for TAPNet
if 'TAPNet' in self.model:
# load optical flow <prev_frame, cur_frame>
optflow = self.load_optflow(filename)
# generate attention map
if abs_idx % utils.num_frames_video:
if self.mf:
# calculate attention map using previous prediction and Motion Flow
attmap = cal_attmap_np(self.attmaps[abs_idx - 1], optflow)
else:
# calculate attention map using previous prediction
attmap = self.attmaps[abs_idx - 1]
'''
EXPERIMENT: uncomment to use
don't use optical flow, directly use previous frame prediction in last epoch
'''
# attmap = self.attmaps[abs_idx - 1]
else:
# first frame of every video, simply use prediction in last epoch without motion flow
attmap = self.attmaps[abs_idx]
# input absolute index and attention map for attention map update
input_dict['abs_idx'] = abs_idx
input_dict['attmap'] = torch.from_numpy(np.expand_dims(attmap,
0)).float()
# gts
mask = self.load_mask(filename, self.mask_folder)
# augment
data = {'image': image, 'mask': mask}
augmented = self.transform(**data)
image, mask = augmented["image"], augmented["mask"]
# input image
input_dict['input'] = img_to_tensor(image)
if self.mode == 'eval':
# in evaluation mode should input the filename of input image
input_dict['input_filename'] = str(filename)
else:
# input gt
input_dict['target'] = torch.from_numpy(mask).long()
return input_dict
def __getitem__(self, index):
self.image_paths = Path(
os.path.join(self.root, self.folder_name[index])).rglob('*.jpg')
image_paths = list(self.image_paths)
index_jump = len(image_paths) / 32
image = []
index_image = 0
while (index_image < len(image_paths)):
x = cv2.imread(str(image_paths[int(
index_image)])) #Image.open(image_paths[int(index_image)])
x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
transformed = self.transform(image=x)
transformed_image = transformed["image"]
image.append(img_to_tensor(transformed_image))
index_image += index_jump
'''
for index_image in range(0, len(image_paths), index_jump):
x = Image.open(image_paths[index_image])
image.append(self.transform(x))
'''
#x = Image.open(self.image_paths[index])
#label = self.json_data[str(self.image_paths[index]).split('/')[-2]+'.mp4']['label']
label = self.json_data[self.folder_name[index] + '.mp4']['label']
if label == 'FAKE':
self.label = torch.tensor(1.0).repeat(32)
else:
self.label = torch.tensor(0.0).repeat(32)
return torch.cat(image).view(-1, 3, 256, 256), self.label
def __getitem__(self, idx):
if self.mode == 'train' and self.equibatch:
group_id = self.group_ids[idx % len(self.group_ids)]
name = random.choice(self.group_names[group_id])
else:
group_id = "unknown"
name = self.names[idx]
pre_img_path = os.path.join(self.data_path, "images", name + "_pre_disaster.png")
post_img_path = os.path.join(self.data_path, "images", name + "_post_disaster.png")
image_pre = cv2.imread(pre_img_path, cv2.IMREAD_COLOR)[:, :, ::-1]
image_post = cv2.imread(post_img_path, cv2.IMREAD_COLOR)[:, :, ::-1]
mask_pre = cv2.imread(os.path.join(self.data_path, "masks", name + "_pre_disaster.png"), cv2.IMREAD_GRAYSCALE)
mask_post = cv2.imread(os.path.join(self.data_path, "masks", name + "_post_disaster.png"), cv2.IMREAD_GRAYSCALE)
rectangles = self.cache.get(self.names[idx], [])
if not rectangles:
self.add_boxes(label(mask_post == 2).astype(np.uint8), rectangles)
if rectangles:
self.cache[self.names[idx]] = rectangles
mask = np.stack([mask_pre, mask_post, mask_post], axis=-1)
sample = self.transforms(image=image_pre, image1=image_post, mask=mask, img_name=name, rectangles=rectangles)
image = np.concatenate([sample['image'], sample['image1']], axis=-1)
sample['img_name'] = name
sample['group_id'] = group_id
mask = np.zeros((5, *sample["mask"].shape[:2]))
for i in range(1, 5):
mask[i - 1, sample["mask"][:, :, 1] == i] = 1
mask[4] = sample["mask"][:, :, 0] / 255
del sample["image1"]
sample['original_mask'] = torch.from_numpy(np.ascontiguousarray(sample["mask"][:, :, 1]))
sample['mask'] = torch.from_numpy(np.ascontiguousarray(mask)).float()
sample['image'] = img_to_tensor(np.ascontiguousarray(image), self.normalize)
return sample
def act(self, state, epsilon):
if random.random() > epsilon:
with torch.no_grad():
state_ = self.Variable(img_to_tensor(state).unsqueeze(0))
q_value = self.forward(state_)
return q_value.max(1)[1].data[0]
else:
return random.randrange(self.num_actions)
def __getitem__(self, idx):
img_file_name = self.file_names[idx]
image = load_image(img_file_name)
if self.mode == 'train':
mask_path = re.sub('.tif', 'segcls.tif', img_file_name)
mask = load_mask(mask_path)
data = {"image": image, "mask": mask}
augmented = self.transform(**data)
image, mask = augmented["image"], augmented["mask"]
return img_to_tensor(image), torch.from_numpy(np.expand_dims(mask, 0)).float()
else:
data = {"image": image}
augmented = self.transform(**data)
image = augmented["image"]
return img_to_tensor(image), str(img_file_name)
def __getitem__(self, idx):
img_file_name = self.file_names[idx]
image = load_image(img_file_name)
mask = load_mask(img_file_name, self.problem_type)
data = {"image": image, "mask": mask}
augmented = self.transform(**data)
image, mask = augmented["image"], augmented["mask"]
if self.mode == 'train':
if self.problem_type == 'binary':
return img_to_tensor(image), torch.from_numpy(
np.expand_dims(mask, 0)).float()
else:
return img_to_tensor(image), torch.from_numpy(mask).long()
else:
return img_to_tensor(image), str(img_file_name)
def image_loader(image):
image = np.asarray(image)
"""load image, returns cuda tensor"""
image = loader(image=image)['image']
image = img_to_tensor(image)
image = torch.autograd.Variable(image, requires_grad=True)
return image.unsqueeze(0)
def standardize(img_soft, img_withbone, threshold, mask=None, vthreshold=None):
if mask is not None:
# lung_org = img_blend(image_org, img_withbone, mask)
lung_soft = img_soft * mask
lung_withbone = img_withbone * mask
else:
lung_soft = img_soft
lung_withbone = img_withbone
if vthreshold is None:
lung_soft_min = lung_soft[lung_soft > 0].min()
lung_soft_max = lung_soft.max()
lung_withbone_min = lung_withbone[lung_withbone > 0].min()
lung_withbone_max = lung_withbone.max()
#vthreshold = (lung_org_min, lung_org_max)
else:
lung_soft_min = vthreshold[0]
lung_soft_max = vthreshold[1]
lung_withbone_min = vthreshold[0]
lung_withbone_max = vthreshold[1]
# 对target进行归范化处理
a_org = (threshold[1] - threshold[0]) / (lung_soft_max - lung_soft_min)
b_org = (threshold[0] * lung_soft_max -
lung_soft_min * threshold[1]) / (lung_soft_max - lung_soft_min)
image_soft = img_soft * a_org + b_org
image_soft[image_soft < threshold[0]] = threshold[0]
image_soft[image_soft > threshold[1]] = threshold[1]
# 对人造原片进行归范化处理
withbone_a_org = (threshold[1] - threshold[0]) / (lung_withbone_max -
lung_withbone_min)
withbone_b_org = (threshold[0] * lung_soft_max - lung_withbone_min *
threshold[1]) / (lung_withbone_max - lung_withbone_min)
image_withbone = img_withbone * withbone_a_org + withbone_b_org
image_withbone[image_soft < threshold[0]] = threshold[0]
image_withbone[image_soft > threshold[1]] = threshold[1]
return img_to_tensor(image_withbone), img_to_tensor(image_soft)
def direct_val(imgs,size):
#img 输入为RGB顺序
imgs = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in imgs]
transforms = create_val_transforms(size)
normalize = {"mean": [0.485, 0.456, 0.406],
"std": [0.229, 0.224, 0.225]}
imgs = [img_to_tensor(transforms(image=each)['image'],normalize).unsqueeze(0) for each in imgs]
imgs = torch.cat(imgs)
return imgs
def direct_val(imgs,size):
#img 输入为RGB顺序
imgs = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in imgs]
stream = create_val_transforms(size)
normalize = {"mean": [0.485, 0.456, 0.406],
"std": [0.229, 0.224, 0.225]}
imgs = [img_to_tensor(stream(each,return_torch=False).data[0],normalize).unsqueeze(0) for each in imgs]
imgs = torch.cat(imgs)
return imgs
def __getitem__(self, idx):
img_file_name = self.file_names[idx]
image = load_image(img_file_name)
if self.mode == 'train':
name_root = img_file_name.split('_')[-1].split('.')[0]
mask_path = os.path.join(
self.mask_dir, 'mask_' + name_root + '.png')
mask = load_mask(mask_path)
data = {"image": image, "mask": mask}
augmented = self.transform(**data)
image, mask = augmented["image"], augmented["mask"]
return img_to_tensor(image), torch.from_numpy(mask).to(torch.long)
else:
data = {"image": image}
augmented = self.transform(**data)
image = augmented["image"]
return img_to_tensor(image), str(img_file_name)
def test_file():
model = Network().cuda().eval()
fn1 = 'pytorch-unflow-master/images/first.png'
fn2 = 'pytorch-unflow-master/images/second.png'
first = img_to_tensor(cv2.cvtColor(cv2.imread(str(fn1)),
cv2.COLOR_BGR2RGB))
second = img_to_tensor(
cv2.cvtColor(cv2.imread(str(fn2)), cv2.COLOR_BGR2RGB))
output = estimate(model, first.unsqueeze(0), second.unsqueeze(0)).squeeze()
objectOutput = open(str('out.flo'), 'wb')
np.array([80, 73, 69, 72], np.uint8).tofile(objectOutput)
np.array([output.size(2), output.size(1)], np.int32).tofile(objectOutput)
np.array(output.numpy().transpose(1, 2, 0),
np.float32).tofile(objectOutput)
out = output.numpy().transpose(1, 2, 0)
print(out.shape, out[0])
objectOutput.close()
return out
def __getitem__(self, index):
if self.args.backend == "pil":
if self.args.grayscale:
image = Image.open(self.image_paths[index]).convert("L")
image = np.array(image)
image = np.expand_dims(image, -1)
else:
image = Image.open(self.image_paths[index])
image = np.array(image)
elif self.args.backend == "cv2":
if self.args.grayscale:
image = cv2.imread(self.image_paths[index],
cv2.IMREAD_GRAYSCALE)
image = np.expand_dims(image, -1)
else:
image = cv2.imread(self.image_paths[index])
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
else:
raise Exception("Backend not implemented")
# Applying Albumentation augments
if self.transforms:
image = self.transforms(image=image)['image']
# Normalize and to Tensor
if self.args.grayscale:
image = img_to_tensor(image)
else:
image = img_to_tensor(image, {
"mean": [0.485, 0.456, 0.406],
"std": [0.229, 0.224, 0.225]
})
# image = torch.from_numpy(image.transpose((2, 0, 1)))
# image = np.transpose(image, (2,0,1)).astype(np.float32)
if self.is_test:
return image
else:
label = self.labels[index]
return image, label #torch.tensor(img, dtype=torch.float32),\
def __getitem__(self, index):
image = self.images[index]
if self.transforms:
image = self.transforms(image=image)['image']
if self.use_masking:
mask_num = random.randrange(1, args.max_mask_num+1)
if self.args.masking_type == 'soft':
mask_prob = 0.25
rand_val = random.random()
if rand_val < mask_prob:
image = time_mask(image, num_masks=mask_num)
rand_val = random.random()
if rand_val < mask_prob:
image = signal_mask(image, num_masks=mask_num)
elif self.args.masking_type == 'hard':
mask_prob = 0.3
rand_val = random.random()
if rand_val < mask_prob:
image = signal_mask(time_mask(image, num_masks=mask_num), num_masks=mask_num)
else:
rand_val = random.random()
if rand_val < 0.3:
image = time_mask(image, num_masks=mask_num)
rand_val = random.random()
if rand_val < 0.3:
image = signal_mask(image, num_masks=mask_num)
data = {}
if self.use_meta:
data['image'] = img_to_tensor(image)
data['meta'] = torch.tensor(self.meta[index]).float()
else:
data['image'] = img_to_tensor(image)
if self.is_test:
return data
else:
label = self.labels[index]
return data, label
def __getitem__(self, idx):
name = self.names[idx]
pre_img_path = os.path.join(self.data_path, "images", "test_pre_" + name + ".png")
post_img_path = os.path.join(self.data_path, "images", "test_post_" + name + ".png")
image_pre = cv2.imread(pre_img_path, cv2.IMREAD_COLOR)[:, :, ::-1]
image_post = cv2.imread(post_img_path, cv2.IMREAD_COLOR)[:, :, ::-1]
image = np.concatenate([image_pre, image_post], axis=-1)
sample = {}
sample['img_name'] = name
sample['image'] = img_to_tensor(np.ascontiguousarray(image), self.normalize)
return sample
def load_sample(self, img_path):
try:
image = cv2.imread(img_path, cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
if self.transforms:
data = self.transforms(image=image)
image = data["image"]
image = img_to_tensor(image, self.normalize)
return image
except:
self.lost.append(img_path)
#pdb.set_trace()
return torch.randn((3, 380, 380))
