167 lines
6.9 KiB
Python
167 lines
6.9 KiB
Python
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
|