Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| import pathlib | |
| from typing import Union, Optional, List, Tuple, Dict, Text, BinaryIO | |
| from PIL import Image | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from .seq_aligner import get_word_inds | |
| def text_under_image(image: np.ndarray, | |
| text: str, | |
| text_color: Tuple[int, int, int] = (0, 0, 0)) -> np.ndarray: | |
| h, w, c = image.shape | |
| offset = int(h * .2) | |
| img = np.ones((h + offset, w, c), dtype=np.uint8) * 255 | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| img[:h] = image | |
| textsize = cv2.getTextSize(text, font, 1, 2)[0] | |
| text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2 | |
| cv2.putText(img, text, (text_x, text_y), font, 1, text_color, 2) | |
| return img | |
| def view_images( | |
| images: Union[np.ndarray, List[np.ndarray]], | |
| num_rows: int = 1, | |
| offset_ratio: float = 0.02, | |
| save_image: bool = False, | |
| fp: Union[Text, pathlib.Path, BinaryIO] = None, | |
| ) -> np.ndarray: | |
| if save_image: | |
| assert fp is not None | |
| if isinstance(images, list): | |
| images = np.concatenate(images, axis=0) | |
| if isinstance(images, np.ndarray) and images.ndim == 4: | |
| num_empty = images.shape[0] % num_rows | |
| else: | |
| images = [images] if not isinstance(images, list) else images | |
| num_empty = len(images) % num_rows | |
| empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255 | |
| images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty | |
| num_items = len(images) | |
| # Calculate the composite image | |
| h, w, c = images[0].shape | |
| offset = int(h * offset_ratio) | |
| num_cols = int(np.ceil(num_items / num_rows)) # count the number of columns | |
| image_h = h * num_rows + offset * (num_rows - 1) | |
| image_w = w * num_cols + offset * (num_cols - 1) | |
| assert image_h > 0, "Invalid image height: {} (num_rows={}, offset_ratio={}, num_items={})".format( | |
| image_h, num_rows, offset_ratio, num_items) | |
| assert image_w > 0, "Invalid image width: {} (num_cols={}, offset_ratio={}, num_items={})".format( | |
| image_w, num_cols, offset_ratio, num_items) | |
| image_ = np.ones((image_h, image_w, 3), dtype=np.uint8) * 255 | |
| # Ensure that the last row is filled with empty images if necessary | |
| if len(images) % num_cols > 0: | |
| empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255 | |
| num_empty = num_cols - len(images) % num_cols | |
| images += [empty_images] * num_empty | |
| for i in range(num_rows): | |
| for j in range(num_cols): | |
| k = i * num_cols + j | |
| if k >= num_items: | |
| break | |
| image_[i * (h + offset): i * (h + offset) + h, j * (w + offset): j * (w + offset) + w] = images[k] | |
| pil_img = Image.fromarray(image_) | |
| if save_image: | |
| pil_img.save(fp) | |
| return pil_img | |
| def update_alpha_time_word(alpha, | |
| bounds: Union[float, Tuple[float, float]], | |
| prompt_ind: int, | |
| word_inds: Optional[torch.Tensor] = None): | |
| if isinstance(bounds, float): | |
| bounds = 0, bounds | |
| start, end = int(bounds[0] * alpha.shape[0]), int(bounds[1] * alpha.shape[0]) | |
| if word_inds is None: | |
| word_inds = torch.arange(alpha.shape[2]) | |
| alpha[: start, prompt_ind, word_inds] = 0 | |
| alpha[start: end, prompt_ind, word_inds] = 1 | |
| alpha[end:, prompt_ind, word_inds] = 0 | |
| return alpha | |
| def get_time_words_attention_alpha(prompts, num_steps, | |
| cross_replace_steps: Union[float, Dict[str, Tuple[float, float]]], | |
| tokenizer, | |
| max_num_words=77): | |
| if type(cross_replace_steps) is not dict: | |
| cross_replace_steps = {"default_": cross_replace_steps} | |
| if "default_" not in cross_replace_steps: | |
| cross_replace_steps["default_"] = (0., 1.) | |
| alpha_time_words = torch.zeros(num_steps + 1, len(prompts) - 1, max_num_words) | |
| for i in range(len(prompts) - 1): | |
| alpha_time_words = update_alpha_time_word(alpha_time_words, cross_replace_steps["default_"], | |
| i) | |
| for key, item in cross_replace_steps.items(): | |
| if key != "default_": | |
| inds = [get_word_inds(prompts[i], key, tokenizer) for i in range(1, len(prompts))] | |
| for i, ind in enumerate(inds): | |
| if len(ind) > 0: | |
| alpha_time_words = update_alpha_time_word(alpha_time_words, item, i, ind) | |
| alpha_time_words = alpha_time_words.reshape(num_steps + 1, len(prompts) - 1, 1, 1, max_num_words) | |
| return alpha_time_words | |