PKUWilliamYang commited on
Commit
25a8215
1 Parent(s): 880ad85

Update vtoonify_model.py

Browse files
Files changed (1) hide show
  1. vtoonify_model.py +4 -0
vtoonify_model.py CHANGED
@@ -148,6 +148,7 @@ class Model():
148
 
149
  video_cap = cv2.VideoCapture(video)
150
  if video_cap.get(7) == 0:
 
151
  return np.zeros((256,256,3), np.uint8), torch.zeros(1,18,512).to(self.device), 'Error: fail to load the video.'
152
  success, frame = video_cap.read()
153
  frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
@@ -158,6 +159,9 @@ class Model():
158
  message = 'Error: no face detected! Please retry or change the video.'
159
  instyle = torch.zeros(1,18,512).to(self.device)
160
  video_cap = cv2.VideoCapture(video)
 
 
 
161
  num = min(300, int(video_cap.get(7)))
162
  if self.device == 'cpu':
163
  num = min(100, num)
 
148
 
149
  video_cap = cv2.VideoCapture(video)
150
  if video_cap.get(7) == 0:
151
+ video_cap.release()
152
  return np.zeros((256,256,3), np.uint8), torch.zeros(1,18,512).to(self.device), 'Error: fail to load the video.'
153
  success, frame = video_cap.read()
154
  frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 
159
  message = 'Error: no face detected! Please retry or change the video.'
160
  instyle = torch.zeros(1,18,512).to(self.device)
161
  video_cap = cv2.VideoCapture(video)
162
+ if video_cap.get(7) == 0:
163
+ video_cap.release()
164
+ return 'default.mp4' instyle, 'Error: fail to load the video.'
165
  num = min(300, int(video_cap.get(7)))
166
  if self.device == 'cpu':
167
  num = min(100, num)