Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| from PIL import Image | |
| import numpy as np | |
| def rescale_image(img, scale, nearest=32, max_size=1280): | |
| w, h = img.size | |
| new_w = int(w * scale) | |
| new_h = int(h * scale) | |
| if new_w > max_size or new_h > max_size: | |
| # Calculate new size keeping aspect ratio | |
| scale = min(max_size / new_w, max_size / new_h) | |
| new_w = int(new_w * scale) | |
| new_h = int(new_h * scale) | |
| # Adjust to nearest multiple | |
| new_w = (new_w // nearest) * nearest | |
| new_h = (new_h // nearest) * nearest | |
| return img.resize((new_w, new_h), Image.LANCZOS), new_w, new_h | |
| def replace_transparent(input_image, target_color): | |
| """ | |
| Converts transparent pixels in a PIL Image to white. | |
| """ | |
| if input_image.mode in ('RGBA', 'LA') or \ | |
| (input_image.mode == 'P' and 'transparency' in input_image.info): | |
| alpha = input_image.convert('RGBA').split()[-1] | |
| bg = Image.new("RGB", input_image.size, target_color) | |
| converted_image = Image.composite(input_image.convert("RGB"), bg, alpha) | |
| return converted_image | |
| else: | |
| # If the image doesn't have an alpha channel, return it as is | |
| return input_image | |
| def padding_image(images, new_width, new_height): | |
| new_image = Image.new('RGB', (new_width, new_height), (255, 255, 255)) | |
| aspect_ratio = images.width / images.height | |
| if new_width / new_height > 1: | |
| if aspect_ratio > new_width / new_height: | |
| new_img_width = new_width | |
| new_img_height = int(new_img_width / aspect_ratio) | |
| else: | |
| new_img_height = new_height | |
| new_img_width = int(new_img_height * aspect_ratio) | |
| else: | |
| if aspect_ratio > new_width / new_height: | |
| new_img_width = new_width | |
| new_img_height = int(new_img_width / aspect_ratio) | |
| else: | |
| new_img_height = new_height | |
| new_img_width = int(new_img_height * aspect_ratio) | |
| resized_img = images.resize((new_img_width, new_img_height)) | |
| paste_x = (new_width - new_img_width) // 2 | |
| paste_y = (new_height - new_img_height) // 2 | |
| new_image.paste(resized_img, (paste_x, paste_y)) | |
| return new_image | |
| def get_image_latent(ref_image=None, sample_size=None, padding=False): | |
| if ref_image is not None: | |
| if isinstance(ref_image, str): | |
| ref_image = Image.open(ref_image).convert("RGB") | |
| if padding: | |
| ref_image = padding_image( | |
| ref_image, sample_size[1], sample_size[0]) | |
| ref_image = ref_image.resize((sample_size[1], sample_size[0])) | |
| ref_image = torch.from_numpy(np.array(ref_image)) | |
| ref_image = ref_image.unsqueeze(0).permute( | |
| [3, 0, 1, 2]).unsqueeze(0) / 255 | |
| elif isinstance(ref_image, Image.Image): | |
| ref_image = ref_image.convert("RGB") | |
| if padding: | |
| ref_image = padding_image( | |
| ref_image, sample_size[1], sample_size[0]) | |
| ref_image = ref_image.resize((sample_size[1], sample_size[0])) | |
| ref_image = torch.from_numpy(np.array(ref_image)) | |
| ref_image = ref_image.unsqueeze(0).permute( | |
| [3, 0, 1, 2]).unsqueeze(0) / 255 | |
| else: | |
| ref_image = torch.from_numpy(np.array(ref_image)) | |
| ref_image = ref_image.unsqueeze(0).permute( | |
| [3, 0, 1, 2]).unsqueeze(0) / 255 | |
| return ref_image |