hugging-apps/echo-memory

1import imageio, os2import numpy as np3from PIL import Image4from tqdm import tqdm5 6 7class LowMemoryVideo:8    def __init__(self, file_name):9        self.reader = imageio.get_reader(file_name)10    11    def __len__(self):12        return self.reader.count_frames()13 14    def __getitem__(self, item):15        return Image.fromarray(np.array(self.reader.get_data(item))).convert("RGB")16 17    def __del__(self):18        self.reader.close()19 20 21def split_file_name(file_name):22    result = []23    number = -124    for i in file_name:25        if ord(i)>=ord("0") and ord(i)<=ord("9"):26            if number == -1:27                number = 028            number = number*10 + ord(i) - ord("0")29        else:30            if number != -1:31                result.append(number)32                number = -133            result.append(i)34    if number != -1:35        result.append(number)36    result = tuple(result)37    return result38 39 40def search_for_images(folder):41    file_list = [i for i in os.listdir(folder) if i.endswith(".jpg") or i.endswith(".png")]42    file_list = [(split_file_name(file_name), file_name) for file_name in file_list]43    file_list = [i[1] for i in sorted(file_list)]44    file_list = [os.path.join(folder, i) for i in file_list]45    return file_list46 47 48class LowMemoryImageFolder:49    def __init__(self, folder, file_list=None):50        if file_list is None:51            self.file_list = search_for_images(folder)52        else:53            self.file_list = [os.path.join(folder, file_name) for file_name in file_list]54    55    def __len__(self):56        return len(self.file_list)57 58    def __getitem__(self, item):59        return Image.open(self.file_list[item]).convert("RGB")60 61    def __del__(self):62        pass63 64 65def crop_and_resize(image, height, width):66    image = np.array(image)67    image_height, image_width, _ = image.shape68    if image_height / image_width < height / width:69        croped_width = int(image_height / height * width)70        left = (image_width - croped_width) // 271        image = image[:, left: left+croped_width]72        image = Image.fromarray(image).resize((width, height))73    else:74        croped_height = int(image_width / width * height)75        left = (image_height - croped_height) // 276        image = image[left: left+croped_height, :]77        image = Image.fromarray(image).resize((width, height))78    return image79 80 81class VideoData:82    def __init__(self, video_file=None, image_folder=None, height=None, width=None, **kwargs):83        if video_file is not None:84            self.data_type = "video"85            self.data = LowMemoryVideo(video_file, **kwargs)86        elif image_folder is not None:87            self.data_type = "images"88            self.data = LowMemoryImageFolder(image_folder, **kwargs)89        else:90            raise ValueError("Cannot open video or image folder")91        self.length = None92        self.set_shape(height, width)93 94    def raw_data(self):95        frames = []96        for i in range(self.__len__()):97            frames.append(self.__getitem__(i))98        return frames99 100    def set_length(self, length):101        self.length = length102 103    def set_shape(self, height, width):104        self.height = height105        self.width = width106 107    def __len__(self):108        if self.length is None:109            return len(self.data)110        else:111            return self.length112 113    def shape(self):114        if self.height is not None and self.width is not None:115            return self.height, self.width116        else:117            height, width, _ = self.__getitem__(0).shape118            return height, width119 120    def __getitem__(self, item):121        frame = self.data.__getitem__(item)122        width, height = frame.size123        if self.height is not None and self.width is not None:124            if self.height != height or self.width != width:125                frame = crop_and_resize(frame, self.height, self.width)126        return frame127 128    def __del__(self):129        pass130 131    def save_images(self, folder):132        os.makedirs(folder, exist_ok=True)133        for i in tqdm(range(self.__len__()), desc="Saving images"):134            frame = self.__getitem__(i)135            frame.save(os.path.join(folder, f"{i}.png"))136 137 138def save_video(frames, save_path, fps, quality=9, ffmpeg_params=None):139    writer = imageio.get_writer(save_path, fps=fps, quality=quality, ffmpeg_params=ffmpeg_params)140    for frame in tqdm(frames, desc="Saving video"):141        frame = np.array(frame)142        writer.append_data(frame)143    writer.close()144 145def save_frames(frames, save_path):146    os.makedirs(save_path, exist_ok=True)147    for i, frame in enumerate(tqdm(frames, desc="Saving images")):148        frame.save(os.path.join(save_path, f"{i}.png"))149