77 lines
1.9 KiB
Python
77 lines
1.9 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from torchvision import transforms
|
|
|
|
from . import net_microAST
|
|
|
|
|
|
class MODEL_MICROAST:
|
|
def __init__(
|
|
self,
|
|
content_dec_path: str,
|
|
style_enc_path: str,
|
|
modulator_path: str,
|
|
decoder_path: str,
|
|
device,
|
|
) -> None:
|
|
content_encoder = net_microAST.Encoder()
|
|
style_encoder = net_microAST.Encoder()
|
|
modulator = net_microAST.Modulator()
|
|
decoder = net_microAST.Decoder()
|
|
|
|
content_encoder.eval()
|
|
style_encoder.eval()
|
|
modulator.eval()
|
|
decoder.eval()
|
|
|
|
content_encoder.load_state_dict(torch.load(content_dec_path))
|
|
style_encoder.load_state_dict(torch.load(style_enc_path))
|
|
modulator.load_state_dict(torch.load(modulator_path))
|
|
decoder.load_state_dict(torch.load(decoder_path))
|
|
|
|
self.network = net_microAST.TestNet(
|
|
content_encoder, style_encoder, modulator, decoder
|
|
)
|
|
self.network.to(device)
|
|
|
|
|
|
def inference_microast(
|
|
src_img: torch.Tensor,
|
|
style_img: torch.Tensor,
|
|
size: int,
|
|
do_crop: bool,
|
|
alpha: float,
|
|
device,
|
|
network,
|
|
):
|
|
content_tf = test_transform(size, do_crop)
|
|
|
|
content = content_tf(src_img).to(device)
|
|
style = content_tf(style_img).to(device)
|
|
style = F.interpolate(
|
|
style,
|
|
size=(content.shape[2], content.shape[3]),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
|
|
# print(f"Resized: {content.shape=}")
|
|
# print(f"Resized: {style.shape=}")
|
|
|
|
torch.cuda.synchronize()
|
|
output = network(content, style, alpha)
|
|
torch.cuda.synchronize()
|
|
|
|
return output
|
|
|
|
|
|
def test_transform(size, crop):
|
|
transform_list = []
|
|
if size != 0:
|
|
transform_list.append(transforms.Resize(size))
|
|
if crop:
|
|
transform_list.append(transforms.CenterCrop(size))
|
|
transform = transforms.Compose(transform_list)
|
|
# [B, C, H, W]
|
|
return transform
|