dream2589632147 commited on
Commit
3a27073
·
verified ·
1 Parent(s): 791da29

Fix GHOST 2.0 parsing mask dimensions

Browse files
Files changed (1) hide show
  1. 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)