Spaces:
Runtime error
Runtime error
File size: 2,138 Bytes
dffe42a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import os
from PIL import Image
class ImageStitcher:
def __init__(
self,
tmp_dir: str,
img_res: int = 64,
img_per_line: int = 10,
verbose: bool = False
):
self.update_tmp_dir(tmp_dir)
self.img_res = img_res if isinstance(img_res, tuple) else (img_res, img_res)
self.img_per_line = img_per_line
self.verbose = verbose
self.reset()
def reset(self):
self.cached_img = None
self.img_num = 0
self.num_lines = 1
self.total_width = self.img_res[0] * self.img_per_line
self.total_height = self.img_res[1] * self.num_lines
def update_tmp_dir(self, tmp_dir: str):
tmp_dir = os.path.abspath(tmp_dir)
os.makedirs(tmp_dir, exist_ok=True)
self.tmp_img_path_template = os.path.join(tmp_dir, "img_%03d.png")
def add(self, img: Image, text: str = None):
img = img.resize(self.img_res)
if self.cached_img is None:
new_img = Image.new('RGBA', (self.total_width, self.total_height))
new_img.paste(img, (0, 0))
else:
num_lines = self.img_num // self.img_per_line + 1
if num_lines > self.num_lines:
self.num_lines = num_lines
self.total_height = self.img_res[1] * self.num_lines
new_img = Image.new('RGBA', (self.total_width, self.total_height))
new_img.paste(self.cached_img, (0, 0))
elif num_lines == self.num_lines:
new_img = self.cached_img
y_offset = self.img_res[1] * (num_lines - 1)
x_offset = self.img_res[0] * (self.img_num % self.img_per_line)
new_img.paste(img, (x_offset, y_offset))
save_path = self.tmp_img_path_template % self.img_num
if text is not None:
save_path = save_path.replace(".png", f"_{text}.png")
new_img.save(save_path)
self.cached_img = new_img
self.img_num += 1
if self.verbose:
print(f"Saved image to {save_path}")
return save_path
|