PKUWilliamYang commited on
Commit
7ab3ef3
1 Parent(s): bf6a25e

Update vtoonify_model.py

Browse files
Files changed (1) hide show
  1. vtoonify_model.py +4 -2
vtoonify_model.py CHANGED
@@ -201,7 +201,8 @@ class Model():
201
 
202
  return 'input.mp4', instyle, 'Successfully rescale the video to (%d, %d)'%(bottom-top, right-left)
203
 
204
- def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float) -> tuple[np.ndarray, str]:
 
205
  if instyle is None or aligned_face is None:
206
  return np.zeros((256,256,3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
207
  if exstyle is None:
@@ -224,7 +225,8 @@ class Model():
224
  print('*** Toonify %dx%d image'%(y_tilde.shape[2], y_tilde.shape[3]))
225
  return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s'%(self.style_name)
226
 
227
- def video_tooniy(self, aligned_video: str, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float) -> tuple[str, str]:
 
228
  if aligned_video is None:
229
  return 'default.mp4', 'Opps, something wrong with the input. Please go to Step 2 and Rescale Video again.'
230
  video_cap = cv2.VideoCapture(aligned_video)
 
201
 
202
  return 'input.mp4', instyle, 'Successfully rescale the video to (%d, %d)'%(bottom-top, right-left)
203
 
204
+ def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple[np.ndarray, str]:
205
+ print(style_type + ' ' + self.style_name)
206
  if instyle is None or aligned_face is None:
207
  return np.zeros((256,256,3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
208
  if exstyle is None:
 
225
  print('*** Toonify %dx%d image'%(y_tilde.shape[2], y_tilde.shape[3]))
226
  return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s'%(self.style_name)
227
 
228
+ def video_tooniy(self, aligned_video: str, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple[str, str]:
229
+ print(style_type + ' ' + self.style_name)
230
  if aligned_video is None:
231
  return 'default.mp4', 'Opps, something wrong with the input. Please go to Step 2 and Rescale Video again.'
232
  video_cap = cv2.VideoCapture(aligned_video)