Spaces:
Running
Running
File size: 4,190 Bytes
966ae59 |
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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
from typing import Dict
import torch
import torch.nn as nn
from torch.nn import functional as nnf
import torchvision
import numpy as np
from scipy.spatial import Delaunay
from shapely.geometry import Point
from shapely.geometry.polygon import Polygon
class ToneLoss(nn.Module):
def __init__(self, cfg):
super(ToneLoss, self).__init__()
self.dist_loss_weight = cfg.dist_loss_weight
self.im_init = None
self.mse_loss = nn.MSELoss()
self.blur = torchvision.transforms.GaussianBlur(
kernel_size=(cfg.pixel_dist_kernel_blur,
cfg.pixel_dist_kernel_blur),
sigma=(cfg.pixel_dist_sigma, cfg.pixel_dist_sigma)
)
self.init_blurred = None
def set_image_init(self, im_init):
self.im_init = im_init
self.init_blurred = self.blur(self.im_init)
def get_scheduler(self, step=None):
if step is not None:
return self.dist_loss_weight * np.exp(-(1 / 5) * ((step - 300) / (20)) ** 2)
else:
return self.dist_loss_weight
def forward(self, cur_raster, step=None):
blurred_cur = self.blur(cur_raster)
return self.mse_loss(self.init_blurred.detach(), blurred_cur) * self.get_scheduler(step)
class ConformalLoss:
def __init__(self, parameters, shape_groups, target_letter: str, device: torch.device):
self.parameters = parameters
self.device = device
self.target_letter = target_letter
self.shape_groups = shape_groups
self.faces = self.init_faces(device)
self.faces_roll_a = [torch.roll(self.faces[i], 1, 1) for i in range(len(self.faces))]
with torch.no_grad():
self.angles = []
self.reset(device)
def get_angles(self, points: torch.Tensor) -> torch.Tensor:
angles_ = []
for i in range(len(self.faces)):
triangles = points[self.faces[i]]
triangles_roll_a = points[self.faces_roll_a[i]]
edges = triangles_roll_a - triangles
length = edges.norm(dim=-1)
edges = edges / (length + 1e-1)[:, :, None]
edges_roll = torch.roll(edges, 1, 1)
cosine = torch.einsum('ned,ned->ne', edges, edges_roll)
angles = torch.arccos(cosine)
angles_.append(angles)
return angles_
def get_letter_inds(self, letter_to_insert):
for group, l in zip(self.shape_groups, self.target_letter):
if l == letter_to_insert:
letter_inds = group.shape_ids
return letter_inds[0], letter_inds[-1], len(letter_inds)
def reset(self, device):
points = torch.cat([point.to(device) for point in self.parameters])
self.angles = self.get_angles(points)
def init_faces(self, device: torch.device) -> torch.tensor:
faces_ = []
for j, c in enumerate(self.target_letter):
points_np = [
self.parameters[i].clone().detach().cpu().numpy()
for i in range(len(self.parameters))
]
start_ind, end_ind, shapes_per_letter = self.get_letter_inds(c)
print(c, "start_ind: ", start_ind.item(), ", end_ind: ", end_ind.item())
holes = []
if shapes_per_letter > 1:
holes = points_np[start_ind + 1:end_ind]
poly = Polygon(points_np[start_ind], holes=holes)
poly = poly.buffer(0)
points_np = np.concatenate(points_np)
faces = Delaunay(points_np).simplices
is_intersect = np.array([poly.contains(Point(points_np[face].mean(0))) for face in faces], dtype=bool)
faces_.append(torch.from_numpy(faces[is_intersect]).to(device, dtype=torch.int64))
return faces_
def __call__(self) -> torch.Tensor:
loss_angles = 0
points = torch.cat(self.parameters).to(self.device)
angles = self.get_angles(points)
for i in range(len(self.faces)):
loss_angles += (nnf.mse_loss(angles[i], self.angles[i]))
return loss_angles
|