Spaces:
Running
Running
Update vtoonify_model.py
Browse files- vtoonify_model.py +10 -9
vtoonify_model.py
CHANGED
@@ -142,15 +142,16 @@ class Model():
|
|
142 |
return frame, instyle, message
|
143 |
|
144 |
#@torch.inference_mode()
|
145 |
-
def detect_and_align_image(self,
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
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:
|
|
|
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:
|