forked from lizhengwei1992/Semantic_Human_Matting
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_camera.py
130 lines (90 loc) · 3.07 KB
/
test_camera.py
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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
'''
test camera
Author: Zhengwei Li
Date : 2018/12/28
'''
import time
import cv2
import torch
import argparse
import numpy as np
import os
import torch.nn.functional as F
parser = argparse.ArgumentParser(description='human matting')
parser.add_argument('--model', default='./', help='preTrained model')
parser.add_argument('--size', type=int, default=256, help='input size')
parser.add_argument('--without_gpu', action='store_true', default=False, help='no use gpu')
args = parser.parse_args()
torch.set_grad_enabled(False)
#################################
#----------------
if args.without_gpu:
print("use CPU !")
device = torch.device('cpu')
else:
if torch.cuda.is_available():
n_gpu = torch.cuda.device_count()
print("----------------------------------------------------------")
print("| use GPU ! || Available GPU number is {} ! |".format(n_gpu))
print("----------------------------------------------------------")
device = torch.device('cuda:0,1')
#################################
#---------------
def load_model(args):
print('Loading model from {}...'.format(args.model))
if args.without_gpu:
myModel = torch.load(args.model, map_location=lambda storage, loc: storage)
else:
myModel = torch.load(args.model)
myModel.eval()
myModel.to(device)
return myModel
def seg_process(args, image, net):
# opencv
origin_h, origin_w, c = image.shape
image_resize = cv2.resize(image, (args.size,args.size), interpolation=cv2.INTER_CUBIC)
image_resize = (image_resize - (104., 112., 121.,)) / 255.0
tensor_4D = torch.FloatTensor(1, 3, args.size, args.size)
tensor_4D[0,:,:,:] = torch.FloatTensor(image_resize.transpose(2,0,1))
inputs = tensor_4D.to(device)
t0 = time.time()
trimap, alpha = net(inputs)
print((time.time() - t0))
if args.without_gpu:
alpha_np = alpha[0,0,:,:].data.numpy()
else:
alpha_np = alpha[0,0,:,:].cpu().data.numpy()
alpha_np = cv2.resize(alpha_np, (origin_w, origin_h), interpolation=cv2.INTER_CUBIC)
fg = np.multiply(alpha_np[..., np.newaxis], image)
bg = image
bg_gray = np.multiply(1-alpha_np[..., np.newaxis], image)
bg_gray = cv2.cvtColor(bg_gray, cv2.COLOR_BGR2GRAY)
bg[:,:,0] = bg_gray
bg[:,:,1] = bg_gray
bg[:,:,2] = bg_gray
# fg[fg<=0] = 0
# fg[fg>255] = 255
# fg = fg.astype(np.uint8)
# out = cv2.addWeighted(fg, 0.7, bg, 0.3, 0)
out = fg + bg
out[out<0] = 0
out[out>255] = 255
out = out.astype(np.uint8)
return out
def camera_seg(args, net):
videoCapture = cv2.VideoCapture(0)
while(1):
# get a frame
ret, frame = videoCapture.read()
frame = cv2.flip(frame,1)
frame_seg = seg_process(args, frame, net)
# show a frame
cv2.imshow("capture", frame_seg)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
videoCapture.release()
def main(args):
myModel = load_model(args)
camera_seg(args, myModel)
if __name__ == "__main__":
main(args)