Update vtoonify_model.py
Browse files- 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,
|
119 |
-
|
120 |
-
|
121 |
-
|
122 |
-
|
123 |
-
|
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'
|
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}")
|