def inference(args):
with open(args.config) as cfg:
config = yaml.load(cfg)
model = get_generator(config['model'])
model.load_state_dict(torch.load(args.weights)['model'])
model = model.cuda()
os.makedirs(os.path.dirname(args.output), exist_ok=True)
img_transforms = transforms.Compose(
[transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
size_transform = Compose([PadIfNeeded(736, 1280)])
crop = CenterCrop(720, 1280)
img = cv2.imread(args.input)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_s = size_transform(image=img)['image']
img_tensor = torch.from_numpy(
np.transpose(img_s / 255, (2, 0, 1)).astype('float32'))
img_tensor = img_transforms(img_tensor)
with torch.no_grad():
img_tensor = Variable(img_tensor.unsqueeze(0).cuda())
result_image = model(img_tensor)
result_image = result_image[0].cpu().float().numpy()
result_image = (np.transpose(result_image, (1, 2, 0)) + 1) / 2.0 * 255.0
result_image = crop(image=result_image)['image']
result_image = result_image.astype('uint8')
cv2.imwrite(args.output, cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR))
def __init__(self, weights, model_name):
super().__init__()
parameters = {'g_name': model_name, 'norm_layer': 'instance'}
self.impl = get_generator(parameters)
checkpoint = torch.load(weights, map_location='cpu')['model']
self.impl.load_state_dict(checkpoint)
self.impl.train(True)
remove_all_batch_norm(self.impl)
def __init__(self, weights_path: str, model_name: str = ''):
with open('config/config.yaml') as cfg:
config = yaml.load(cfg)
model = get_generator(model_name or config['model'])
model.load_state_dict(torch.load(weights_path)['model'])
self.model = model.cuda()
self.model.train(True)
# GAN inference should be in train mode to use actual stats in norm layers,
# it's not a bug
self.normalize_fn = get_normalize()
def __init__(self, weights_path, model_name=''):
with open('config/config.yaml') as cfg:
config = yaml.load(cfg)
model = get_generator(model_name or config['model'])
model.load_state_dict(
torch.load(weights_path,
map_location=lambda storage, loc: storage)['model'])
if torch.cuda.is_available():
self.model = model.cuda()
else:
self.model = model
self.model.train(True)
# GAN inference should be in train mode to use actual stats in norm layers,
# it's not a bug
self.normalize_fn = get_normalize()
def __init__(self,
weights_path: str = "fpn_inception.h5",
model_name: str = 'fpn_inception'):
with open('config/config.yaml') as cfg:
config = yaml.load(cfg)
self.use_gpu = torch.cuda.is_available()
model = get_generator(model_name, config)
model.load_state_dict(
torch.load(weights_path, map_location="cpu")['model'])
if self.use_gpu:
self.model = model.cuda()
else:
self.model = model
self.model.train(True)
# GAN inference should be in train mode to use actual stats in norm layers,
# it's not a bug
self.normalize_fn = get_normalize()
def build(self):
print("build is called")
networks.get_generator()
import torch
import torchvision
import torch.nn as nn
import torch
import torch.nn as nn
from pretrainedmodels import inceptionresnetv2
from torchsummary import summary
import torch.nn.functional as F
from torch.jit import ScriptModule, script_method, trace
import functools
from torchsummary import summary
with open('config/config.yaml') as cfg:
config = yaml.load(cfg, Loader=yaml.FullLoader)
model = get_generator(config['model'])
weights_path = 'fpn_inception.h5'
model.load_state_dict(torch.load(weights_path)['model'])
model.eval()
model = model.cpu()
# summary(model, input_size=(3, 256, 256))
# model = model.cuda()
model.train(True)
#model = model.cpu()
example = torch.rand(1, 3, 128, 128)# .to('cuda')
traced_script_module = torch.jit.trace(model.module, example)
traced_script_module.save("model128.pt")