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('--------------------------------')