100 lines
2.7 KiB
Python
100 lines
2.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torchvision.transforms as transforms
|
|
|
|
from .models import Transformer
|
|
|
|
|
|
class video_Style_transfer(nn.Module):
|
|
def __init__(self, ckpt_path: str, encoder_path: str, decoder_path: str):
|
|
super(video_Style_transfer, self).__init__()
|
|
self.ckpt_path = ckpt_path
|
|
self.model = Transformer(encoder_path, decoder_path)
|
|
self.load_weights()
|
|
|
|
def load_weights(self):
|
|
# print("Loading model from checkpoint")
|
|
ckpt = torch.load(self.ckpt_path, map_location="cpu")
|
|
self.model.load_state_dict(get_keys(ckpt, "model"), strict=False)
|
|
|
|
def forward(
|
|
self,
|
|
content_frames,
|
|
style_images,
|
|
content_type="image",
|
|
id_loss="transfer",
|
|
tab=None,
|
|
):
|
|
transfer_result = self.model.forward(
|
|
content_frames, style_images, content_type, id_loss, tab
|
|
)
|
|
return transfer_result
|
|
|
|
|
|
def inference_unist(
|
|
src: torch.Tensor,
|
|
src_style: torch.Tensor,
|
|
device,
|
|
size: int,
|
|
do_crop: bool,
|
|
network: video_Style_transfer,
|
|
content_type="image",
|
|
):
|
|
content_tf = test_transform(size, do_crop)
|
|
content = content_tf(src).unsqueeze(0).to(device) # [1, B, C, H, W]
|
|
|
|
style = content_tf(src_style).to(device) # [B, C, H, W]
|
|
|
|
if content_type == "video":
|
|
style = match_shape(style, content[0])
|
|
|
|
y_hat = network.forward(content, style, content_type, tab="inference")
|
|
output = un_normalize_batch(y_hat) # [B, C, H, W]
|
|
|
|
return output
|
|
|
|
|
|
def test_transform(size, crop, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
|
|
transform_list = []
|
|
|
|
if crop and size != 0:
|
|
transform_list.append(transforms.Resize(size))
|
|
transform_list.append(transforms.CenterCrop(size))
|
|
elif size != 0:
|
|
transform_list.append(transforms.Resize((size, size)))
|
|
|
|
transform_list.append(transforms.Normalize(mean, std))
|
|
|
|
transform = transforms.Compose(transform_list)
|
|
# [B, C, H, W]
|
|
return transform
|
|
|
|
|
|
def un_normalize_batch(tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
|
|
for t, m, s in zip(tensor, mean, std):
|
|
t.mul_(s).add_(m)
|
|
return tensor
|
|
|
|
|
|
def match_shape(a: torch.Tensor, b: torch.Tensor):
|
|
B_a, C, H, W = a.shape
|
|
B_b = b.shape[0]
|
|
|
|
if B_a == B_b:
|
|
return a
|
|
elif B_a == 1:
|
|
return a.expand(B_b, -1, -1, -1)
|
|
elif B_a < B_b:
|
|
repeat_factor = B_b // B_a
|
|
remainder = B_b % B_a
|
|
return torch.cat([a.repeat(repeat_factor, 1, 1, 1), a[:remainder]], dim=0)
|
|
else: # B_a > B_b
|
|
return a[:B_b]
|
|
|
|
|
|
def get_keys(d, name):
|
|
if "state_dict" in d:
|
|
d = d["state_dict"]
|
|
d_filt = {k[len(name) + 1 :]: v for k, v in d.items() if k[: len(name)] == name}
|
|
return d_filt
|