File size: 3,970 Bytes
210e8a2 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | #!/usr/bin/python
# -*- encoding: utf-8 -*-
import numpy as np
from model import BiSeNet
import torch
import os
import os.path as osp
from PIL import Image
import torchvision.transforms as transforms
import cv2
from pathlib import Path
import configargparse
import tqdm
# import ttach as tta
def vis_parsing_maps(im, parsing_anno, stride, save_im=False, save_path='vis_results/parsing_map_on_im.jpg',
img_size=(512, 512)):
im = np.array(im)
vis_im = im.copy().astype(np.uint8)
vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
vis_parsing_anno = cv2.resize(
vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST)
vis_parsing_anno_color = np.zeros(
(vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + np.array([255, 255, 255]) # + 255
num_of_class = np.max(vis_parsing_anno)
# print(num_of_class)
for pi in range(1, 14):
index = np.where(vis_parsing_anno == pi)
vis_parsing_anno_color[index[0], index[1], :] = np.array([255, 0, 0])
for pi in [11]:
index = np.where(vis_parsing_anno == pi)
vis_parsing_anno_color[index[0], index[1], :] = np.array([100, 100, 100])
for pi in range(14, 16):
index = np.where(vis_parsing_anno == pi)
vis_parsing_anno_color[index[0], index[1], :] = np.array([0, 255, 0])
for pi in range(16, 17):
index = np.where(vis_parsing_anno == pi)
vis_parsing_anno_color[index[0], index[1], :] = np.array([0, 0, 255])
for pi in (17, 18):
index = np.where(vis_parsing_anno == pi)
vis_parsing_anno_color[index[0], index[1], :] = np.array([0, 0, 0])
for pi in range(18, num_of_class+1):
index = np.where(vis_parsing_anno == pi)
vis_parsing_anno_color[index[0], index[1], :] = np.array([255, 0, 0])
vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
index = np.where(vis_parsing_anno == num_of_class-1)
vis_im = cv2.resize(vis_parsing_anno_color, img_size,
interpolation=cv2.INTER_NEAREST)
if save_im:
cv2.imwrite(save_path, vis_im)
def evaluate(respth='./res/test_res', dspth='./data', cp='model_final_diss.pth'):
Path(respth).mkdir(parents=True, exist_ok=True)
print(f'[INFO] loading model...')
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.load_state_dict(torch.load(cp))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
image_paths = sorted(os.listdir(dspth))
with torch.no_grad():
for image_path in tqdm.tqdm(image_paths):
if image_path.endswith('.jpg') or image_path.endswith('.png'):
img = Image.open(osp.join(dspth, image_path))
ori_size = img.size
image = img.resize((512, 512), Image.BILINEAR)
image = image.convert("RGB")
img = to_tensor(image)
# test-time augmentation.
inputs = torch.unsqueeze(img, 0) # [1, 3, 512, 512]
outputs = net(inputs.cuda())
parsing = outputs.mean(0).cpu().numpy().argmax(0)
image_path = int(image_path[:-4])
image_path = str(image_path) + '.png'
vis_parsing_maps(image, parsing, stride=1, save_im=True, save_path=osp.join(respth, image_path), img_size=ori_size)
if __name__ == "__main__":
parser = configargparse.ArgumentParser()
parser.add_argument('--respath', type=str, default='./result/', help='result path for label')
parser.add_argument('--imgpath', type=str, default='./imgs/', help='path for input images')
parser.add_argument('--modelpath', type=str, default='data_utils/face_parsing/79999_iter.pth')
args = parser.parse_args()
evaluate(respth=args.respath, dspth=args.imgpath, cp=args.modelpath)
|