AgainstEntropy/Kanji-Streaming
0
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 