Spaces:
Runtime error
Runtime error
File size: 1,749 Bytes
82ef366 |
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 |
import pathlib
from typing import Optional, Tuple
import numpy as np
from PIL import Image
from config import RunConfig
def load_images(cfg: RunConfig, save_path: Optional[pathlib.Path] = None) -> Tuple[Image.Image, Image.Image]:
image_style = load_size(cfg.app_image_path)
image_struct = load_size(cfg.struct_image_path)
return image_style, image_struct
def load_size(image_path: pathlib.Path,
left: int = 0,
right: int = 0,
top: int = 0,
bottom: int = 0,
size: int = 512) -> Image.Image:
if type(image_path) is str or type(image_path) is pathlib.PosixPath:
image = np.array(Image.open(image_path).convert('RGB'))
else:
image = image_path
h, w, c = image.shape
left = min(left, w - 1)
right = min(right, w - left - 1)
top = min(top, h - left - 1)
bottom = min(bottom, h - top - 1)
image = image[top:h - bottom, left:w - right]
h, w, c = image.shape
if h < w:
offset = (w - h) // 2
image = image[:, offset:offset + h]
elif w < h:
offset = (h - w) // 2
image = image[offset:offset + w]
image = np.array(Image.fromarray(image).resize((size, size)))
return image
def save_generated_masks(model, cfg: RunConfig):
tensor2im(model.image_app_mask_32).save(cfg.output_path / f"mask_style_32.png")
tensor2im(model.image_struct_mask_32).save(cfg.output_path / f"mask_struct_32.png")
tensor2im(model.image_app_mask_64).save(cfg.output_path / f"mask_style_64.png")
tensor2im(model.image_struct_mask_64).save(cfg.output_path / f"mask_struct_64.png")
def tensor2im(x) -> Image.Image:
return Image.fromarray(x.cpu().numpy().astype(np.uint8) * 255) |