Spaces:
Running on Zero
Running on Zero
Fix GHOST 2.0 parsing mask dimensions
Browse files- ghost_runtime.py +1 -1
ghost_runtime.py
CHANGED
|
@@ -172,7 +172,7 @@ class Ghost2Engine:
|
|
| 172 |
mask = torch.zeros_like(parsing, dtype=torch.float32)
|
| 173 |
for index in range(1, 21):
|
| 174 |
mask[parsing == index] = 1.0
|
| 175 |
-
return mask[0]
|
| 176 |
|
| 177 |
def _process(self, frame: np.ndarray, target: bool = False):
|
| 178 |
faces = self.detector.get(frame)
|
|
|
|
| 172 |
mask = torch.zeros_like(parsing, dtype=torch.float32)
|
| 173 |
for index in range(1, 21):
|
| 174 |
mask[parsing == index] = 1.0
|
| 175 |
+
return mask[0, 0] if mask.ndim == 4 else mask[0]
|
| 176 |
|
| 177 |
def _process(self, frame: np.ndarray, target: bool = False):
|
| 178 |
faces = self.detector.get(frame)
|