| 45 | |
| 46 | |
| 47 | class VideoDataset_Motion(Dataset): |
| 48 | |
| 49 | def __init__( |
| 50 | self, |
| 51 | config, |
| 52 | download_folder_path, |
| 53 | csv_relative_path, |
| 54 | video_relative_path, |
| 55 | is_diy_test = False, |
| 56 | ) -> None: |
| 57 | super().__init__() |
| 58 | |
| 59 | # Gen Size Settings |
| 60 | # self.height_range = config["height_range"] |
| 61 | # self.max_aspect_ratio = config["max_aspect_ratio"] |
| 62 | self.target_height = config["target_height"] |
| 63 | self.target_width = config["target_width"] |
| 64 | self.sample_accelerate_factor = config["sample_accelerate_factor"] |
| 65 | self.train_frame_num_range = config["train_frame_num_range"] |
| 66 | |
| 67 | # Condition Settings (Text, Motion, etc.) |
| 68 | self.empty_text_prompt = config["empty_text_prompt"] |
| 69 | self.dot_radius = int(config["dot_radius"]) |
| 70 | self.point_keep_ratio = config["point_keep_ratio"] # Point selection mechanism |
| 71 | self.faster_motion_prob = config["faster_motion_prob"] |
| 72 | |
| 73 | # Other Settings |
| 74 | self.download_folder_path = download_folder_path |
| 75 | self.is_diy_test = is_diy_test |
| 76 | self.config = config |
| 77 | self.video_folder_path = os.path.join(download_folder_path, video_relative_path) |
| 78 | csv_folder_path = os.path.join(download_folder_path, csv_relative_path) |
| 79 | |
| 80 | |
| 81 | # Sanity Check |
| 82 | assert(os.path.exists(csv_folder_path)) |
| 83 | assert(self.point_keep_ratio <= 1.0) |
| 84 | |
| 85 | |
| 86 | |
| 87 | # Read the CSV files |
| 88 | info_lists = [] |
| 89 | for csv_file_name in os.listdir(csv_folder_path): # Read all csv files |
| 90 | csv_file_path = os.path.join(csv_folder_path, csv_file_name) |
| 91 | |
| 92 | with open(csv_file_path) as file_obj: |
| 93 | reader_obj = csv.reader(file_obj) |
| 94 | |
| 95 | # Iterate over each row in the csv |
| 96 | for idx, row in enumerate(reader_obj): |
| 97 | if idx == 0: |
| 98 | elements = dict() |
| 99 | for element_idx, key in enumerate(row): |
| 100 | elements[key] = element_idx |
| 101 | continue |
| 102 | |
| 103 | # Read the important information |
| 104 | info_lists.append(row) |