import cv2 import pandas as pd import gc import lpips from PIL import Image, ImageOps import requests import torch from torch import nn from torch.nn import functional as F import torchvision import torchvision.transforms as T import torchvision.transforms.functional as TF from tqdm import tqdm from resize_right import resize from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults import numpy as np from numpy import asarray from .settings import DiscoDiffusionSettings from . import disco_utils class MakeCutouts(nn.Module): def __init__(self, cut_size, cutn, skip_augs=False): super().__init__() self.cut_size = cut_size self.cutn = cutn self.skip_augs = skip_augs self.augs = T.Compose([ T.RandomHorizontalFlip(p=0.5), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomAffine(degrees=15, translate=(0.1, 0.1)), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomPerspective(distortion_scale=0.4, p=0.7), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomGrayscale(p=0.15), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), # T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1), ]) def forward(self, input): input = T.Pad(input.shape[2]//4, fill=0)(input) sideY, sideX = input.shape[2:4] max_size = min(sideX, sideY) cutouts = [] for ch in range(self.cutn): if ch > self.cutn - self.cutn//4: cutout = input.clone() else: size = int(max_size * torch.zeros(1,).normal_(mean=.8, std=.3).clip(float(self.cut_size/max_size), 1.)) offsetx = torch.randint(0, abs(sideX - size + 1), ()) offsety = torch.randint(0, abs(sideY - size + 1), ()) cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size] if not self.skip_augs: cutout = self.augs(cutout) cutouts.append(disco_utils.resample(cutout, (self.cut_size, self.cut_size))) del cutout cutouts = torch.cat(cutouts, dim=0) return cutouts cutout_debug = False padargs = {} class MakeCutoutsDango(nn.Module): def __init__(self, animation_mode: str, skip_augs, cut_size, Overview=4, InnerCrop = 0, IC_Size_Pow=0.5, IC_Grey_P = 0.2 ): super().__init__() self.cut_size = cut_size self.skip_augs = skip_augs self.Overview = Overview self.InnerCrop = InnerCrop self.IC_Size_Pow = IC_Size_Pow self.IC_Grey_P = IC_Grey_P if animation_mode == 'None': self.augs = T.Compose([ T.RandomHorizontalFlip(p=0.5), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomAffine(degrees=10, translate=(0.05, 0.05), interpolation = T.InterpolationMode.BILINEAR), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomGrayscale(p=0.1), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1), ]) elif animation_mode == 'Video Input': self.augs = T.Compose([ T.RandomHorizontalFlip(p=0.5), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomAffine(degrees=15, translate=(0.1, 0.1)), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomPerspective(distortion_scale=0.4, p=0.7), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomGrayscale(p=0.15), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), # T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1), ]) elif animation_mode == '2D' or animation_mode == '3D': self.augs = T.Compose([ T.RandomHorizontalFlip(p=0.4), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomAffine(degrees=10, translate=(0.05, 0.05), interpolation = T.InterpolationMode.BILINEAR), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.RandomGrayscale(p=0.1), T.Lambda(lambda x: x + torch.randn_like(x) * 0.01), T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.3), ]) def forward(self, input): cutouts = [] gray = T.Grayscale(3) sideY, sideX = input.shape[2:4] max_size = min(sideX, sideY) min_size = min(sideX, sideY, self.cut_size) max(sideX, sideY) output_shape = [1,3,self.cut_size,self.cut_size] [1,3,self.cut_size+2,self.cut_size+2] pad_input = F.pad(input,((sideY-max_size)//2,(sideY-max_size)//2,(sideX-max_size)//2,(sideX-max_size)//2), **padargs) cutout = resize(pad_input, out_shape=output_shape) if self.Overview>0: if self.Overview<=4: if self.Overview>=1: cutouts.append(cutout) if self.Overview>=2: cutouts.append(gray(cutout)) if self.Overview>=3: cutouts.append(TF.hflip(cutout)) if self.Overview==4: cutouts.append(gray(TF.hflip(cutout))) else: cutout = resize(pad_input, out_shape=output_shape) for _ in range(self.Overview): cutouts.append(cutout) if cutout_debug: # if is_colab: # TF.to_pil_image(cutouts[0].clamp(0, 1).squeeze(0)).save("/content/cutout_overview0.jpg",quality=99) TF.to_pil_image(cutouts[0].clamp(0, 1).squeeze(0)).save("cutout_overview0.jpg",quality=99) if self.InnerCrop >0: for i in range(self.InnerCrop): size = int(torch.rand([])**self.IC_Size_Pow * (max_size - min_size) + min_size) offsetx = torch.randint(0, sideX - size + 1, ()) offsety = torch.randint(0, sideY - size + 1, ()) cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size] if i <= int(self.IC_Grey_P * self.InnerCrop): cutout = gray(cutout) cutout = resize(cutout, out_shape=output_shape) cutouts.append(cutout) if cutout_debug: # if is_colab: # TF.to_pil_image(cutouts[-1].clamp(0, 1).squeeze(0)).save("/content/cutout_InnerCrop.jpg",quality=99) # else: TF.to_pil_image(cutouts[-1].clamp(0, 1).squeeze(0)).save("cutout_InnerCrop.jpg",quality=99) cutouts = torch.cat(cutouts) if self.skip_augs is not True: cutouts=self.augs(cutouts) return cutouts