FlexTheAi/Flexstorydiff
0
1from email.mime import image2import torch3import base644import gradio as gr5import numpy as np6from PIL import Image,ImageOps,ImageDraw, ImageFont7from io import BytesIO8import random9MAX_COLORS = 1210def get_random_bool():11 return random.choice([True, False])12 13def add_white_border(input_image, border_width=10):14 """15 为PIL图像添加指定宽度的白色边框。16 17 :param input_image: PIL图像对象18 :param border_width: 边框宽度（单位：像素）19 :return: 带有白色边框的PIL图像对象20 """21 border_color = 'white' # 白色边框22 # 添加边框23 img_with_border = ImageOps.expand(input_image, border=border_width, fill=border_color)24 return img_with_border25 26def process_mulline_text(draw, text, font, max_width):27 """28 Draw the text on an image with word wrapping.29 """30 lines = [] # Store the lines of text here31 words = text.split()32 33 # Start building lines of text, and wrap when necessary34 current_line = ""35 for word in words:36 test_line = f"{current_line} {word}".strip()37 # Check the width of the line with this word added38 bbox = draw.textbbox((0, 0), test_line, font=font)39 text_left, text_top, text_right, text_bottom = bbox40 41 width, _ = (text_right - text_left, text_bottom - text_top)42 43 if width <= max_width:44 # If it fits, add this word to the current line45 current_line = test_line46 else:47 # If not, store the line and start a new one48 lines.append(current_line)49 current_line = word50 # Add the last line51 lines.append(current_line)52 return lines 53 54 55 56def add_caption(image, text, position = "bottom-mid", font = None, text_color= 'black', bg_color = (255, 255, 255) , bg_opacity = 200):57 if text == "":58 return image59 image = image.convert("RGBA")60 draw = ImageDraw.Draw(image)61 width, height = image.size62 lines = process_mulline_text(draw,text,font,width)63 text_positions = []64 maxwidth = 065 for ind, line in enumerate(lines[::-1]):66 bbox = draw.textbbox((0, 0), line, font=font)67 text_left, text_top, text_right, text_bottom = bbox68 text_width, text_height = (text_right - text_left, text_bottom - text_top)69 if position == 'bottom-right':70 text_position = (width - text_width - 10, height - (text_height + 20))71 elif position == 'bottom-left':72 text_position = (10, height - (text_height + 20))73 elif position == 'bottom-mid':74 text_position = ((width - text_width) // 2, height - (text_height + 20) ) # 居中文本75 height = text_position[1]76 maxwidth = max(maxwidth,text_width)77 text_positions.append(text_position)78 rectpos = (width - maxwidth) // 279 rectangle_position = [rectpos - 5, text_positions[-1][1] - 5, rectpos + maxwidth + 5, text_positions[0][1] + text_height + 5]80 image_with_transparency = Image.new('RGBA', image.size)81 draw_with_transparency = ImageDraw.Draw(image_with_transparency)82 draw_with_transparency.rectangle(rectangle_position, fill=bg_color + (bg_opacity,))83 84 image.paste(Image.alpha_composite(image.convert('RGBA'), image_with_transparency))85 print(ind,text_position)86 draw = ImageDraw.Draw(image)87 for ind, line in enumerate(lines[::-1]):88 text_position = text_positions[ind]89 draw.text(text_position, line, fill=text_color, font=font)90 91 return image.convert('RGB')92 93def get_comic(images,types = "4panel",captions = [],font = None,pad_image = None):94 if pad_image == None:95 pad_image = Image.open("./images/pad_images.png")96 97 if types == "No typesetting (default)":98 return images99 elif types == "Four Pannel":100 return get_comic_4panel(images,captions,font,pad_image)101 else: # "Classic Comic Style"102 return get_comic_classical(images,captions,font,pad_image)103 104def get_caption_group(images_groups,captions = []):105 caption_groups = []106 for i in range(len(images_groups)):107 length = len(images_groups[i])108 caption_groups.append(captions[:length])109 captions = captions[length:]110 if len(caption_groups[-1]) < len(images_groups[-1]):111 caption_groups[-1] = caption_groups[-1] + [""] * (len(images_groups[-1]) - len(caption_groups[-1]))112 return caption_groups113 114def get_comic_classical(images,captions = None,font = None,pad_image = None):115 if pad_image == None:116 raise ValueError("pad_image is None")117 images = [add_white_border(image) for image in images]118 pad_image = pad_image.resize(images[0].size, Image.LANCZOS)119 images_groups = distribute_images2(images,pad_image)120 print(images_groups)121 if captions != None:122 captions_groups = get_caption_group(images_groups,captions)123 # print(images_groups)124 row_images = []125 for ind, img_group in enumerate(images_groups):126 row_images.append(get_row_image2(img_group ,captions= captions_groups[ind] if captions != None else None,font = font)) 127 128 return [combine_images_vertically_with_resize(row_images)]129 130 131 132def get_comic_4panel(images,captions = [],font = None,pad_image = None):133 if pad_image == None:134 raise ValueError("pad_image is None")135 pad_image = pad_image.resize(images[0].size, Image.LANCZOS)136 images = [add_white_border(image) for image in images]137 assert len(captions) == len(images)138 for i,caption in enumerate(captions):139 images[i] = add_caption(images[i],caption,font = font)140 images_nums = len(images)141 pad_nums = int((4 - images_nums % 4) % 4) 142 images = images + [pad_image for _ in range(pad_nums)]143 comics = []144 assert len(images)%4 == 0145 for i in range(len(images)//4):146 comics.append(combine_images_vertically_with_resize([combine_images_horizontally(images[i*4:i*4+2]), combine_images_horizontally(images[i*4+2:i*4+4])]))147 148 return comics149 150def get_row_image(images):151 row_image_arr = []152 if len(images)>3:153 stack_img_nums = (len(images) - 2)//2154 else:155 stack_img_nums = 0156 while(len(images)>0):157 if stack_img_nums <=0:158 row_image_arr.append(images[0])159 images = images[1:]160 elif len(images)>stack_img_nums*2:161 if get_random_bool():162 row_image_arr.append(concat_images_vertically_and_scale(images[:2]))163 images = images[2:]164 stack_img_nums -=1165 else:166 row_image_arr.append(images[0])167 images = images[1:]168 else:169 row_image_arr.append(concat_images_vertically_and_scale(images[:2]))170 images = images[2:]171 stack_img_nums-=1172 return combine_images_horizontally(row_image_arr)173 174def get_row_image2(images,captions = None, font = None):175 row_image_arr = []176 if len(images)== 6:177 sequence_list = [1,1,2,2]178 elif len(images)== 4:179 sequence_list = [1,1,2]180 else:181 raise ValueError("images nums is not 4 or 6 found",len(images))182 random.shuffle(sequence_list)183 index = 0184 for length in sequence_list:185 if length == 1:186 if captions != None:187 images_tmp = add_caption(images[0],text = captions[index],font= font)188 else:189 images_tmp = images[0]190 row_image_arr.append( images_tmp)191 images = images[1:]192 index +=1193 elif length == 2:194 row_image_arr.append(concat_images_vertically_and_scale(images[:2]))195 images = images[2:]196 index +=2197 198 return combine_images_horizontally(row_image_arr)199 200 201 202def concat_images_vertically_and_scale(images,scale_factor=2):203 # 加载所有图像204 # 确保所有图像的宽度一致205 widths = [img.width for img in images]206 if not all(width == widths[0] for width in widths):207 raise ValueError('All images must have the same width.')208 209 # 计算总高度210 total_height = sum(img.height for img in images)211 212 # 创建新的图像，宽度与原图相同，高度为所有图像高度之和213 max_width = max(widths)214 concatenated_image = Image.new('RGB', (max_width, total_height))215 216 # 竖直拼接图像217 current_height = 0218 for img in images:219 concatenated_image.paste(img, (0, current_height))220 current_height += img.height221 222 # 缩放图像为1/n高度223 new_height = concatenated_image.height // scale_factor224 new_width = concatenated_image.width // scale_factor225 resized_image = concatenated_image.resize((new_width, new_height), Image.LANCZOS)226 227 return resized_image228 229 230def combine_images_horizontally(images):231 # 读取所有图片并存入列表232 233 # 获取每幅图像的宽度和高度234 widths, heights = zip(*(i.size for i in images))235 236 # 计算总宽度和最大高度237 total_width = sum(widths)238 max_height = max(heights)239 240 # 创建新的空白图片，用于拼接241 new_im = Image.new('RGB', (total_width, max_height))242 243 # 将图片横向拼接244 x_offset = 0245 for im in images:246 new_im.paste(im, (x_offset, 0))247 x_offset += im.width248 249 return new_im250 251def combine_images_vertically_with_resize(images):252 253 # 获取所有图片的宽度和高度254 widths, heights = zip(*(i.size for i in images))255 256 # 确定新图片的宽度，即所有图片中最小的宽度257 min_width = min(widths)258 259 # 调整图片尺寸以保持宽度一致，长宽比不变260 resized_images = []261 for img in images:262 # 计算新高度保持图片长宽比263 new_height = int(min_width * img.height / img.width)264 # 调整图片大小265 resized_img = img.resize((min_width, new_height), Image.LANCZOS)266 resized_images.append(resized_img)267 268 # 计算所有调整尺寸后图片的总高度269 total_height = sum(img.height for img in resized_images)270 271 # 创建一个足够宽和高的新图片对象272 new_im = Image.new('RGB', (min_width, total_height))273 274 # 竖直拼接图片275 y_offset = 0276 for im in resized_images:277 new_im.paste(im, (0, y_offset))278 y_offset += im.height279 280 return new_im281 282def distribute_images2(images, pad_image):283 groups = []284 remaining = len(images)285 if len(images) <= 8:286 group_sizes = [4]287 else:288 group_sizes = [4, 6]289 290 size_index = 0291 while remaining > 0:292 size = group_sizes[size_index%len(group_sizes)] 293 if remaining < size and remaining < min(group_sizes):294 size = min(group_sizes) 295 if remaining > size:296 new_group = images[-remaining: -remaining + size]297 else:298 new_group = images[-remaining:]299 groups.append(new_group)300 size_index += 1301 remaining -= size302 print(remaining,groups)303 groups[-1] = groups[-1] + [pad_image for _ in range(-remaining)]304 305 return groups306 307 308def distribute_images(images, group_sizes=(4, 3, 2)):309 groups = []310 remaining = len(images)311 312 while remaining > 0:313 # 优先分配最大组（4张图片），再考虑3张，最后处理2张314 for size in sorted(group_sizes, reverse=True):315 # 如果剩下的图片数量大于等于当前组大小，或者为图片总数时（也就是第一次迭代）316 # 开始创建新组317 if remaining >= size or remaining == len(images):318 if remaining > size:319 new_group = images[-remaining: -remaining + size]320 else:321 new_group = images[-remaining:]322 groups.append(new_group)323 remaining -= size324 break325 # 如果剩下的图片少于最小的组大小（2张）并且已经有组了，就把剩下的图片加到最后一个组326 elif remaining < min(group_sizes) and groups:327 groups[-1].extend(images[-remaining:])328 remaining = 0329 330 return groups331 332def create_binary_matrix(img_arr, target_color):333 mask = np.all(img_arr == target_color, axis=-1)334 binary_matrix = mask.astype(int)335 return binary_matrix336 337def preprocess_mask(mask_, h, w, device):338 mask = np.array(mask_)339 mask = mask.astype(np.float32)340 mask = mask[None, None]341 mask[mask < 0.5] = 0342 mask[mask >= 0.5] = 1343 mask = torch.from_numpy(mask).to(device)344 mask = torch.nn.functional.interpolate(mask, size=(h, w), mode='nearest')345 return mask346 347def process_sketch(canvas_data):348 binary_matrixes = []349 base64_img = canvas_data['image']350 image_data = base64.b64decode(base64_img.split(',')[1])351 image = Image.open(BytesIO(image_data)).convert("RGB")352 im2arr = np.array(image)353 colors = [tuple(map(int, rgb[4:-1].split(','))) for rgb in canvas_data['colors']]354 colors_fixed = []355 356 r, g, b = 255, 255, 255357 binary_matrix = create_binary_matrix(im2arr, (r,g,b))358 binary_matrixes.append(binary_matrix)359 binary_matrix_ = np.repeat(np.expand_dims(binary_matrix, axis=(-1)), 3, axis=(-1))360 colored_map = binary_matrix_*(r,g,b) + (1-binary_matrix_)*(50,50,50)361 colors_fixed.append(gr.update(value=colored_map.astype(np.uint8)))362 363 for color in colors:364 r, g, b = color365 if any(c != 255 for c in (r, g, b)):366 binary_matrix = create_binary_matrix(im2arr, (r,g,b))367 binary_matrixes.append(binary_matrix)368 binary_matrix_ = np.repeat(np.expand_dims(binary_matrix, axis=(-1)), 3, axis=(-1))369 colored_map = binary_matrix_*(r,g,b) + (1-binary_matrix_)*(50,50,50)370 colors_fixed.append(gr.update(value=colored_map.astype(np.uint8)))371 372 visibilities = []373 colors = []374 for n in range(MAX_COLORS):375 visibilities.append(gr.update(visible=False))376 colors.append(gr.update())377 for n in range(len(colors_fixed)):378 visibilities[n] = gr.update(visible=True)379 colors[n] = colors_fixed[n]380 381 return [gr.update(visible=True), binary_matrixes, *visibilities, *colors]382 383def process_prompts(binary_matrixes, *seg_prompts):384 return [gr.update(visible=True), gr.update(value=' , '.join(seg_prompts[:len(binary_matrixes)]))]385 386def process_example(layout_path, all_prompts, seed_):387 388 all_prompts = all_prompts.split('***')389 390 binary_matrixes = []391 colors_fixed = []392 393 im2arr = np.array(Image.open(layout_path))[:,:,:3]394 unique, counts = np.unique(np.reshape(im2arr,(-1,3)), axis=0, return_counts=True)395 sorted_idx = np.argsort(-counts)396 397 binary_matrix = create_binary_matrix(im2arr, (0,0,0))398 binary_matrixes.append(binary_matrix)399 binary_matrix_ = np.repeat(np.expand_dims(binary_matrix, axis=(-1)), 3, axis=(-1))400 colored_map = binary_matrix_*(255,255,255) + (1-binary_matrix_)*(50,50,50)401 colors_fixed.append(gr.update(value=colored_map.astype(np.uint8)))402 403 for i in range(len(all_prompts)-1):404 r, g, b = unique[sorted_idx[i]]405 if any(c != 255 for c in (r, g, b)) and any(c != 0 for c in (r, g, b)):406 binary_matrix = create_binary_matrix(im2arr, (r,g,b))407 binary_matrixes.append(binary_matrix)408 binary_matrix_ = np.repeat(np.expand_dims(binary_matrix, axis=(-1)), 3, axis=(-1))409 colored_map = binary_matrix_*(r,g,b) + (1-binary_matrix_)*(50,50,50)410 colors_fixed.append(gr.update(value=colored_map.astype(np.uint8)))411 412 visibilities = []413 colors = []414 prompts = []415 for n in range(MAX_COLORS):416 visibilities.append(gr.update(visible=False))417 colors.append(gr.update())418 prompts.append(gr.update())419 420 for n in range(len(colors_fixed)):421 visibilities[n] = gr.update(visible=True)422 colors[n] = colors_fixed[n]423 prompts[n] = all_prompts[n+1]424 425 return [gr.update(visible=True), binary_matrixes, *visibilities, *colors, *prompts,426 gr.update(visible=True), gr.update(value=all_prompts[0]), int(seed_)]