| import imageio, os | |
| import numpy as np | |
| from PIL import Image | |
| def read_video(file_name): | |
| reader = imageio.get_reader(file_name) | |
| video = [] | |
| for frame in reader: | |
| frame = np.array(frame) | |
| video.append(frame) | |
| reader.close() | |
| return video | |
| def get_video_fps(file_name): | |
| reader = imageio.get_reader(file_name) | |
| fps = reader.get_meta_data()["fps"] | |
| reader.close() | |
| return fps | |
| def save_video(frames_path, video_path, num_frames, fps): | |
| writer = imageio.get_writer(video_path, fps=fps, quality=9) | |
| for i in range(num_frames): | |
| frame = np.array(Image.open(os.path.join(frames_path, "%05d.png" % i))) | |
| writer.append_data(frame) | |
| writer.close() | |
| return video_path | |
| class LowMemoryVideo: | |
| def __init__(self, file_name): | |
| self.reader = imageio.get_reader(file_name) | |
| def __len__(self): | |
| return self.reader.count_frames() | |
| def __getitem__(self, item): | |
| return np.array(self.reader.get_data(item)) | |
| def __del__(self): | |
| self.reader.close() | |
| def split_file_name(file_name): | |
| result = [] | |
| number = -1 | |
| for i in file_name: | |
| if ord(i)>=ord("0") and ord(i)<=ord("9"): | |
| if number == -1: | |
| number = 0 | |
| number = number*10 + ord(i) - ord("0") | |
| else: | |
| if number != -1: | |
| result.append(number) | |
| number = -1 | |
| result.append(i) | |
| if number != -1: | |
| result.append(number) | |
| result = tuple(result) | |
| return result | |
| def search_for_images(folder): | |
| file_list = [i for i in os.listdir(folder) if i.endswith(".jpg") or i.endswith(".png")] | |
| file_list = [(split_file_name(file_name), file_name) for file_name in file_list] | |
| file_list = [i[1] for i in sorted(file_list)] | |
| file_list = [os.path.join(folder, i) for i in file_list] | |
| return file_list | |
| def read_images(folder): | |
| file_list = search_for_images(folder) | |
| frames = [np.array(Image.open(i)) for i in file_list] | |
| return frames | |
| class LowMemoryImageFolder: | |
| def __init__(self, folder, file_list=None): | |
| if file_list is None: | |
| self.file_list = search_for_images(folder) | |
| else: | |
| self.file_list = [os.path.join(folder, file_name) for file_name in file_list] | |
| def __len__(self): | |
| return len(self.file_list) | |
| def __getitem__(self, item): | |
| return np.array(Image.open(self.file_list[item])) | |
| def __del__(self): | |
| pass | |
| class VideoData: | |
| def __init__(self, video_file, image_folder, **kwargs): | |
| if video_file is not None: | |
| self.data_type = "video" | |
| self.data = LowMemoryVideo(video_file, **kwargs) | |
| elif image_folder is not None: | |
| self.data_type = "images" | |
| self.data = LowMemoryImageFolder(image_folder, **kwargs) | |
| else: | |
| raise ValueError("Cannot open video or image folder") | |
| self.length = None | |
| self.height = None | |
| self.width = None | |
| def raw_data(self): | |
| frames = [] | |
| for i in range(self.__len__()): | |
| frames.append(self.__getitem__(i)) | |
| return frames | |
| def set_length(self, length): | |
| self.length = length | |
| def set_shape(self, height, width): | |
| self.height = height | |
| self.width = width | |
| def __len__(self): | |
| if self.length is None: | |
| return len(self.data) | |
| else: | |
| return self.length | |
| def shape(self): | |
| if self.height is not None and self.width is not None: | |
| return self.height, self.width | |
| else: | |
| height, width, _ = self.__getitem__(0).shape | |
| return height, width | |
| def __getitem__(self, item): | |
| frame = self.data.__getitem__(item) | |
| height, width, _ = frame.shape | |
| if self.height is not None and self.width is not None: | |
| if self.height != height or self.width != width: | |
| frame = Image.fromarray(frame).resize((self.width, self.height)) | |
| frame = np.array(frame) | |
| return frame | |
| def __del__(self): | |
| pass | |