def __getitem__(self, num_frame):
""" Returns (X_curr_left, X_curr_right), (X_flow_left, X_flow_right), (X_next_left, X_next_right)"""
frame_path_curr_left = self.pic_direcs_list[num_frame][0]
frame_path_curr_right = self.pic_direcs_list[num_frame][1]
frame_curr_left = Image.open(frame_path_curr_left)
frame_curr_right = Image.open(frame_path_curr_right)
frame_curr_left = self.transform(frame_curr_left)
frame_curr_right = self.transform(frame_curr_right)
# Get directory for next picture
frame_name_num = os.path.basename(
frame_path_curr_left) # ie "0001.png"
suffix = os.path.splitext(frame_path_curr_left)[1] # .png
direc_without_num_frame_left = frame_path_curr_left.replace(
frame_name_num, '')
direc_without_num_frame_right = frame_path_curr_right.replace(
frame_name_num, '')
frame_name_string = frame_name_num.replace(suffix, '') # 0001
num_digits = len(frame_name_string) # 0001 -> num_digits = 4
int_curr = int(frame_name_string)
int_next = int_curr + 1 # n+1 as int
string_next = str(int_next)
num_digits_next = len(string_next)
string_next = '0' * (
num_digits - num_digits_next) + string_next + suffix # 0002.png
frame_path_next_left = os.path.join(direc_without_num_frame_left,
string_next)
frame_path_next_right = os.path.join(direc_without_num_frame_right,
string_next)
frame_next_left = Image.open(frame_path_next_left)
frame_next_right = Image.open(frame_path_next_right)
frame_next_left = self.transform(frame_next_left)
frame_next_right = self.transform(frame_next_right)
# Get directory for flow of current picture
direc_without_num_frame_left = os.path.normpath(
direc_without_num_frame_left)
direc_without_left = os.path.split(direc_without_num_frame_left)[0]
scene_name = os.path.split(direc_without_left)[
1] # ie: a_rain_of_stones_x2
flow_name_num_left = "OpticalFlowIntoFuture_" + frame_name_string + "_L.flo" # OpticalFlowIntoFuture_0000_L.flo
flow_name_num_right = "OpticalFlowIntoFuture_" + frame_name_string + "_R.flo" # OpticalFlowIntoFuture_0000_L.flo
flow_path_left = os.path.join(self.flow_path, scene_name, "left",
flow_name_num_left)
flow_path_right = os.path.join(self.flow_path, scene_name, "right",
flow_name_num_right)
flow_left = utils_dataset.readFlow(flow_path_left)
flow_right = utils_dataset.readFlow(flow_path_right)
return (frame_curr_left,
frame_curr_right), (flow_left, flow_right), (frame_next_left,
frame_next_right)
def show_flow_on_image(img_path, flow_path):
# TODO: Optical flow doesn't work correctly (prob dimensions). Fix
img = Image.open(img_path)
transform_to_tensor = transforms.Compose(
[transforms.Resize(image_size),
transforms.ToTensor()])
img = transform_to_tensor(img)
C, H, W = img.shape
im_batch = torch.ones((1, C, H, W))
im_batch[0, :, :, :] = img
img = im_batch
flow = utils_dataset.readFlow(flow_path)
flow = flow[..., ::] - np.zeros_like(flow)
flow = torch.from_numpy(flow)
flow = flow.unsqueeze(0)
new_image, mask = utils.apply_flow(img, flow)
new_image = np.asarray(255 * new_image).astype("uint8")
mask = 255 * np.asarray(mask).astype("uint8")
print((new_image.shape, type(new_image[0, 0, 0])))
print(mask.shape, type(mask))
Image.fromarray(new_image).show()
Image.fromarray(mask).show()
def show_flow_after_style():
model = "models/model_test_temp_1e5.pth"
has_cuda = 1
# img = Image.open("../Data/Monkaa/RGB_cleanpass/a_rain_of_stones_x2/left/0097.png")
img = Image.open("../Data/Monkaa/frames_cleanpass/eating_x2/left/0049.png")
left_frame_stylized = stylize(has_cuda, img, model)
flow = utils_dataset.readFlow(
"../Data/Monkaa/optical_flow/eating_x2/left/OpticalFlowIntoFuture_0049_L.pfm"
)
print("a")
flow = flow[..., ::-1] - np.zeros_like(flow)
flow = torch.from_numpy(flow)
flow = flow.unsqueeze(0)
left_frame_stylized = left_frame_stylized.unsqueeze(0)
new_image, mask = utils.apply_flow(left_frame_stylized, flow)
stylized_frame = new_image.clone().clamp(0, 255).numpy()
stylized_frame = stylized_frame.astype("uint8")
img = Image.fromarray(stylized_frame)
img.show()
def resize_flow_script():
og_dir = "./Data/Monkaa/optical_flow"
new_dir = "./Data/Monkaa/optical_flow_resized"
new_width = 640
new_height = 360
for root, dirs, files in os.walk(og_dir):
for file in files:
if file.endswith(".pfm"):
flow_path = os.path.join(root, file)
flow = utils_dataset.readFlow(flow_path)
flow_resized = resize_flow(flow, new_width, new_height)
dir_folder, lr_folder = os.path.split(root)
_, scene = os.path.split(dir_folder)
flow_resized_path = os.path.join(new_dir, scene, lr_folder)
if not os.path.exists(flow_resized_path):
os.makedirs(flow_resized_path)
file_name = file.replace('.pfm', '.flo')
utils_dataset.write(os.path.join(flow_resized_path, file_name),
flow_resized)
def apply_flow(img, flow_path):
flow = utils_dataset.readFlow(flow_path)
flow = np.round(flow)
height, width, _ = np.asarray(img).shape
new_pixel_place = np.indices((height, width)).transpose(1, 2, 0)
new_pixel_place = new_pixel_place + flow[:, :, ::-1]
new_pixel_place = new_pixel_place.astype(int)
im_array = np.asarray(img)
new_image = np.zeros_like(im_array)
valid_indices = np.where((new_pixel_place[:, :, 0] >= 0)
& (new_pixel_place[:, :, 0] < height)
& (new_pixel_place[:, :, 1] >= 0)
& (new_pixel_place[:, :, 1] < width))
new_pixel_place = new_pixel_place[valid_indices[0], valid_indices[1], :]
new_image[new_pixel_place[:, 0],
new_pixel_place[:, 1], :] = im_array[valid_indices[0],
valid_indices[1], :]
mask = np.zeros_like(img)
mask[new_pixel_place[:, 0], new_pixel_place[:, 1]] = 1
return new_image, mask
flow = flow[yy, xx, :]
flow[:, :, 0] = flow[:, :, 0] / width_ration
flow[:, :, 1] = flow[:, :, 1] / height_ratio
return flow
og_dir = "./Data/Monkaa/optical_flow" # TODO: Check for forward/ backward directories
new_dir = "./Data/Monkaa/optical_flow_resized"
new_height = 256
new_width = 256
for root, dirs, files in os.walk(og_dir):
for file in files:
if file.endswith(".pfm"):
flow_path = os.path.join(root, file)
flow = utils_dataset.readFlow(flow_path)
flow_resized = resize_flow(flow, new_width, new_height)
dir_folder, lr_folder = os.path.split(root)
_, scene = os.path.split(dir_folder)
flow_resized_path = os.path.join(new_dir, scene, lr_folder)
if not os.path.exists(flow_resized_path):
os.makedirs(flow_resized_path)
file_name = file.replace('.pfm', '.flo')
utils_dataset.write(os.path.join(flow_resized_path, file_name),
flow_resized)
# print(root)
# print(dirs)
# print(files)
# print('--------------------------------')