fm.PyTorch/test.py at master · ash11sh/fm.PyTorch · GitHub

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
#!/usr/bin/python
# -*- encoding: utf-8 -*-

import torch
import os
from model import BiSeNet
import os.path as osp
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
import cv2
import torch.nn as nn


def vis_parsing_maps(im, parsing_anno, stride, save_im=False, save_path='vis_results/parsing_map_on_im.jpg'):
	# Colors for all 20 parts
	part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
				   [255, 0, 85], [255, 0, 170],
				   [0, 255, 0], [85, 255, 0], [170, 255, 0],
				   [0, 255, 85], [0, 255, 170],
				   [0, 0, 255], [85, 0, 255], [170, 0, 255],
				   [0, 85, 255], [0, 170, 255],
				   [255, 255, 0], [255, 255, 85], [255, 255, 170],
				   [255, 0, 255], [255, 85, 255], [255, 170, 255],
				   [0, 255, 255], [85, 255, 255], [170, 255, 255]]

	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)) + 255

	num_of_class = np.max(vis_parsing_anno)

	for pi in range(1, num_of_class + 1):
		index = np.where(vis_parsing_anno == pi)
		vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]

	vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
	# print(vis_parsing_anno_color.shape, vis_im.shape)
	vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0)

	# Save result or not
	if save_im:
		cv2.imwrite(save_path[:-4] +'.png', vis_parsing_anno)
		cv2.imwrite(save_path, vis_im, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
	return vis_parsing_anno
	# return vis_im


def evaluate(potrait, cp='cp/79999_iter.pth'):

	# if not os.path.exists(respth):
	#     os.makedirs(respth)

	n_classes = 19
	net = BiSeNet(n_classes=n_classes)
	# net.cuda()
	# use_cuda = torch.cuda.is_available()
	device = torch.device("cpu")

	net.to(device)
	net.load_state_dict(torch.load(cp,map_location='cpu'))
	net.eval()

	to_tensor = transforms.Compose([
		transforms.ToTensor(),
		transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
	])

	with torch.no_grad():
		# img = Image.open(image_path)
		img = Image.fromarray(cv2.cvtColor(potrait, cv2.COLOR_BGR2RGB))
		# image = img.resize((512, 512), Image.BILINEAR)
		img = to_tensor(img)

		img = torch.unsqueeze(img, 0)
		# img = img.cuda()
		out = net(img)[0]
		# parsing = out.squeeze(0).cpu().numpy().argmax(0)
		parsing = out.squeeze(0).numpy().argmax(0)
		# print(parsing)
		# print(np.unique(parsing))

		# vis_parsing_maps(image, parsing, stride=1, save_im=False, save_path=osp.join(respth, dspth))
		return parsing

if __name__ == "__main__":
	evaluate(dspth='/home/zll/data/CelebAMask-HQ/test-img/116.jpg', cp='79999_iter.pth')