CoolFace
Apppublic

AgainstEntropy/Kanji-Streaming

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
image_utils.py68 linesDownload Raw Back to root
1import os2 3from PIL import Image4 5 6class ImageStitcher:7    def __init__(8            self, 9            tmp_dir: str,10            img_res: int = 64,11            img_per_line: int = 10,12            verbose: bool = False13        ):14        self.update_tmp_dir(tmp_dir)15 16        self.img_res = img_res if isinstance(img_res, tuple) else (img_res, img_res)17        self.img_per_line = img_per_line18 19        self.verbose = verbose20 21        self.reset()22 23    def reset(self):24        self.cached_img = None25        self.img_num = 026        self.num_lines = 127 28        self.total_width = self.img_res[0] * self.img_per_line29        self.total_height = self.img_res[1] * self.num_lines30 31    def update_tmp_dir(self, tmp_dir: str):32        tmp_dir = os.path.abspath(tmp_dir)33        os.makedirs(tmp_dir, exist_ok=True)34        self.tmp_img_path_template = os.path.join(tmp_dir, "img_%03d.png")35 36    def add(self, img: Image, text: str = None):37        38        img = img.resize(self.img_res)39 40        if self.cached_img is None:41            new_img = Image.new('RGBA', (self.total_width, self.total_height))42            new_img.paste(img, (0, 0))43        else:44            num_lines = self.img_num // self.img_per_line + 145            if num_lines > self.num_lines:46                self.num_lines = num_lines47                self.total_height = self.img_res[1] * self.num_lines48                new_img = Image.new('RGBA', (self.total_width, self.total_height))49                new_img.paste(self.cached_img, (0, 0))50            elif num_lines == self.num_lines:51                new_img = self.cached_img52 53            y_offset = self.img_res[1] * (num_lines - 1)54            x_offset = self.img_res[0] * (self.img_num % self.img_per_line)55            new_img.paste(img, (x_offset, y_offset))56 57        save_path = self.tmp_img_path_template % self.img_num58        if text is not None:59            save_path = save_path.replace(".png", f"_{text}.png")60        new_img.save(save_path)61        self.cached_img = new_img62        self.img_num += 163 64        if self.verbose:65            print(f"Saved image to {save_path}")66 67        return save_path68