| 45 | import pdb |
| 46 | |
| 47 | def get_input(args): |
| 48 | input_path = args.input_path |
| 49 | transform_video = transforms.Compose([ |
| 50 | video_transforms.ToTensorVideo(), # TCHW |
| 51 | video_transforms.ResizeVideo((args.image_h, args.image_w)), |
| 52 | transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True) |
| 53 | ]) |
| 54 | if input_path is not None: |
| 55 | print(f'loading video from {input_path}') |
| 56 | if os.path.isdir(input_path): |
| 57 | file_list = os.listdir(input_path) |
| 58 | video_frames = [] |
| 59 | if args.mask_type.startswith('onelast'): |
| 60 | num = int(args.mask_type.split('onelast')[-1]) |
| 61 | # get first and last frame |
| 62 | first_frame_path = os.path.join(input_path, natsorted(file_list)[0]) |
| 63 | last_frame_path = os.path.join(input_path, natsorted(file_list)[-1]) |
| 64 | first_frame = torch.as_tensor(np.array(Image.open(first_frame_path), dtype=np.uint8, copy=True)).unsqueeze(0) |
| 65 | last_frame = torch.as_tensor(np.array(Image.open(last_frame_path), dtype=np.uint8, copy=True)).unsqueeze(0) |
| 66 | for i in range(num): |
| 67 | video_frames.append(first_frame) |
| 68 | # add zeros to frames |
| 69 | num_zeros = args.num_frames-2*num |
| 70 | for i in range(num_zeros): |
| 71 | zeros = torch.zeros_like(first_frame) |
| 72 | video_frames.append(zeros) |
| 73 | for i in range(num): |
| 74 | video_frames.append(last_frame) |
| 75 | n = 0 |
| 76 | video_frames = torch.cat(video_frames, dim=0).permute(0, 3, 1, 2) # f,c,h,w |
| 77 | video_frames = transform_video(video_frames) |
| 78 | else: |
| 79 | for file in file_list: |
| 80 | if file.endswith('jpg') or file.endswith('png'): |
| 81 | image = torch.as_tensor(np.array(Image.open(file), dtype=np.uint8, copy=True)).unsqueeze(0) |
| 82 | video_frames.append(image) |
| 83 | else: |
| 84 | continue |
| 85 | n = 0 |
| 86 | video_frames = torch.cat(video_frames, dim=0).permute(0, 3, 1, 2) # f,c,h,w |
| 87 | video_frames = transform_video(video_frames) |
| 88 | return video_frames, n |
| 89 | elif os.path.isfile(input_path): |
| 90 | _, full_file_name = os.path.split(input_path) |
| 91 | file_name, extension = os.path.splitext(full_file_name) |
| 92 | if extension == '.jpg' or extension == '.png': |
| 93 | print("loading the input image") |
| 94 | video_frames = [] |
| 95 | num = int(args.mask_type.split('first')[-1]) |
| 96 | first_frame = torch.as_tensor(np.array(Image.open(input_path), dtype=np.uint8, copy=True)).unsqueeze(0) |
| 97 | for i in range(num): |
| 98 | video_frames.append(first_frame) |
| 99 | num_zeros = args.num_frames-num |
| 100 | for i in range(num_zeros): |
| 101 | zeros = torch.zeros_like(first_frame) |
| 102 | video_frames.append(zeros) |
| 103 | n = 0 |
| 104 | video_frames = torch.cat(video_frames, dim=0).permute(0, 3, 1, 2) # f,c,h,w |