I wanted to test it with 1 image and measure elapsed time. I used suggested well trained pth file. And it takes 1.5 second to run 1 image. Here is the code if someone wants to try it too
create a file called main2.py and add the below code. And then change img_path and model_path, accordingly.
import torch
from torch.backends import cudnn
import cv2
from parameters import *
from trainer import Trainer
from tester import Tester
from data_loader import CustomDataLoader
from utils import make_folder
from augmentations import *
import os.path as osp
import torch
import timeit
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from torchvision.utils import save_image
# from torchvision import transforms
import torchvision.transforms as transforms
from networks import get_model
from utils import *
# from PIL import Image
import time
from metrics import SegMetric
from time import time
facial_names = ['background', 'skin', 'nose', 'eyeglass', 'left_eye', 'right_eye', 'left_brow', 'right_brow',
'left_ear', 'right_ear', 'mouth', 'upper_lip', 'lower_lip', 'hair', 'hat', 'earring',
'necklace',
'neck', 'cloth']
cudnn.enabled = True
cudnn.benchmark = True
cudnn.deterministic = False
torch.cuda.manual_seed(2020)
def load_model(model_path):
model_name="FaceParseNet50"
model= get_model(model_name, n_classes=19, pretrained=False).cuda()
torch_saved_obj_path=(model_path)
torch_saved_obj= torch.load(torch_saved_obj_path )
model.load_state_dict(torch_saved_obj)
model.eval()
return model
def preprocess_img(img):
transform = transforms.ToTensor()
resizer = transforms.Resize(512)
tensor_img = transform(img)
tensor_img=resizer(tensor_img)
tensor_img=torch.unsqueeze(tensor_img, 0)
assert tensor_img.shape[0] == 1
assert tensor_img.shape[1] == 3
assert tensor_img.shape[2] == 512
assert tensor_img.shape[3] == 512
return tensor_img
def face_parser(img, model_name,model_path, color_results):
t0=time()
model=load_model(model_path)
t1 =time()
img=preprocess_img(img)
t2 = time()
img = img.cuda()
torch.cuda.synchronize()
t3 = time()
with torch.no_grad():
outputs = model(img)
t4 = time()
if model_name == "FaceParseNet":
outputs = outputs[0][-1]
h, w = 512, 512
outputs = F.interpolate(outputs, (h, w), mode='bilinear', align_corners=True)
t5 = time()
pred = outputs.data.max(1)[1].cpu().numpy() # Matrix index
t6 = time()
# outputs=model_output_processor(outputs)
if color_results: # Whether color the test results to png files
imsize=512
labels_predict_plain = generate_label_plain(outputs, imsize)
t7 = time()
r=labels_predict_plain[0]
print(type(r))
cv2.imwrite("/home/enes/result.png",r)
t8 = time()
print("model load",t1-t0)
print("preprocess",t2-t1)
print("preprocess2",t3-t2)
print("model output",t4-t3)
print("interpolate",t5-t4)
print("pred",t6-t5)
print("generate_label_plain",t7-t6)
print("imwrite",t8-t7)
return outputs
img_path="/home/enes/lab/faceparser/FaceParsing.PyTorch/Data_preprocessing/CelebA-HQ-img/1.jpg"
img = cv2.imread(img_path)
model_name="FaceParseNet"
model_path="./models/FaceParseNet50/38_G.pth"
face_parser(img,model_name,model_path, color_results=True)
I wanted to test it with 1 image and measure elapsed time. I used suggested well trained pth file. And it takes 1.5 second to run 1 image. Here is the code if someone wants to try it too