File size: 2,219 Bytes
91fc62a |
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 |
import importlib
import PIL
import pytorch_lightning as pl
import torch.utils.data
import wandb
from typing import Union
from torchvision import transforms
from utils_.loss import VGGPerceptualLoss
from utils_.visualization import *
import torch.nn.functional as F
import matplotlib.pyplot as plt
class Model(pl.LightningModule):
def __init__(self, **kwargs):
super().__init__()
self.save_hyperparameters()
self.encoder = importlib.import_module('models.' + self.hparams.encoder).Encoder(self.hparams)
self.decoder = importlib.import_module('models.' + self.hparams.decoder).Decoder(self.hparams)
self.batch_size = self.hparams.batch_size
self.vgg_loss = VGGPerceptualLoss()
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(0.5, 0.5)
])
def forward(self, x: PIL.Image.Image) -> PIL.Image.Image:
"""
:param x: a PIL image
:return: an edge map of the same size as x with values in [0, 1] (normalized by max)
"""
w, h = x.size
x = self.transform(x).unsqueeze(0)
x = x.to(self.device)
kp = self.encoder({'img': x})['keypoints']
edge_map = self.decoder.rasterize(kp, output_size=64)
bs = edge_map.shape[0]
edge_map = edge_map / (1e-8 + edge_map.reshape(bs, 1, -1).max(dim=2, keepdim=True)[0].reshape(bs, 1, 1, 1))
edge_map = torch.cat([edge_map] * 3, dim=1)
edge_map = F.interpolate(edge_map, size=(h, w), mode='bilinear', align_corners=False)
x = torch.clamp(edge_map + (x * 0.5 + 0.5)*0.5, min=0, max=1)
x = transforms.ToPILImage()(x[0].detach().cpu())
fig = plt.figure(figsize=(1, h/w), dpi=w)
fig.tight_layout(pad=0)
plt.axis('off')
plt.imshow(x)
kp = kp[0].detach().cpu() * 0.5 + 0.5
kp[:, 1] *= w
kp[:, 0] *= h
plt.scatter(kp[:, 1], kp[:, 0], s=min(w/h, min(1, h/w)), marker='o')
ncols, nrows = fig.canvas.get_width_height()
fig.canvas.draw()
plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8).reshape(nrows, ncols, 3)
plt.close(fig)
return plot
|