Ashrafb commited on
Commit
2a82e53
·
verified ·
1 Parent(s): 7612bf6

Update vtoonify_model.py

Browse files
Files changed (1) hide show
  1. vtoonify_model.py +10 -8
vtoonify_model.py CHANGED
@@ -115,13 +115,15 @@ class Model():
115
 
116
  return landmarks_68
117
 
118
- def detect_and_align_image(self, frame_bgr: np.ndarray, top: int, bottom: int, left: int, right: int) -> tuple:
119
- if frame_bgr is None:
120
- return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load the image.'
121
-
122
- # Proceed with detection and alignment
123
- return self.detect_and_align(frame_bgr, top, bottom, left, right)
124
-
 
 
125
  def detect_and_align(self, frame, top, bottom, left, right, return_para=False):
126
  message = 'Error: no face detected! Please retry or change the photo.'
127
  instyle = None
@@ -218,7 +220,7 @@ class Model():
218
  logging.info(f"Output from VToonify shape: {y_tilde.shape}")
219
  print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
220
 
221
- 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)
222
 
223
  except Exception as e:
224
  logging.error(f"Error during model execution: {e}")
 
115
 
116
  return landmarks_68
117
 
118
+ def detect_and_align_image(self, image: str, top: int, bottom: int, left: int, right: int
119
+ ) -> tuple[np.ndarray, torch.Tensor, str]:
120
+ if image is None:
121
+ return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load empty file.'
122
+ frame = cv2.imread(image)
123
+ if frame is None:
124
+ return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load the image.'
125
+ frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
126
+ return self.detect_and_align(frame, top, bottom, left, right)
127
  def detect_and_align(self, frame, top, bottom, left, right, return_para=False):
128
  message = 'Error: no face detected! Please retry or change the photo.'
129
  instyle = None
 
220
  logging.info(f"Output from VToonify shape: {y_tilde.shape}")
221
  print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
222
 
223
+ 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)
224
 
225
  except Exception as e:
226
  logging.error(f"Error during model execution: {e}")