Ashrafb commited on
Commit
a7e780b
·
verified ·
1 Parent(s): 92105c7

Update vtoonify_model.py

Browse files
Files changed (1) hide show
  1. vtoonify_model.py +9 -10
vtoonify_model.py CHANGED
@@ -142,16 +142,15 @@ class Model():
142
  return frame, instyle, message
143
 
144
  #@torch.inference_mode()
145
- def detect_and_align_image(self, image: str, top: int, bottom: int, left: int, right: int
146
- ) -> tuple[np.ndarray, torch.Tensor, str]:
147
- if image is None:
148
- return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load empty file.'
149
- frame = cv2.imread(image)
150
- if frame is None:
151
- return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load the image.'
152
- frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
153
- return self.detect_and_align(frame, top, bottom, left, right)
154
-
155
  def detect_and_align_video(self, video: str, top: int, bottom: int, left: int, right: int
156
  ) -> tuple[np.ndarray, torch.Tensor, str]:
157
  if video is None:
 
142
  return frame, instyle, message
143
 
144
  #@torch.inference_mode()
145
+ def detect_and_align_image(self, frame_rgb: np.ndarray, top: int, bottom: int, left: int, right: int) -> tuple[np.ndarray, torch.Tensor, str]:
146
+ if frame_rgb is None:
147
+ return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load the image.'
148
+
149
+ # Convert RGB to BGR
150
+ frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
151
+
152
+ return self.detect_and_align(frame_bgr, top, bottom, left, right)
153
+
 
154
  def detect_and_align_video(self, video: str, top: int, bottom: int, left: int, right: int
155
  ) -> tuple[np.ndarray, torch.Tensor, str]:
156
  if video is None: