debug

app.py

@@ -68,8 +68,7 @@ def main():
                 with gr.Row():
                     label_image = gr.Image(label='Label Image',
                                            type='numpy',
-                                           elem_id='label-image',
-                                           interactive=False)
+                                           elem_id='label-image')
                 with gr.Row():
                     shape_text = gr.Textbox(
                         label='Shape Description',

model.py

@@ -64,6 +64,7 @@ class Model:
         self.config['device'] = device
         self._download_models()
         self.model = SampleFromPoseModel(self.config)
+        self.model.batch_size = 1
 
     def _load_config(self) -> dict:
         path = 'Text2Human/configs/sample_from_pose.yml'
@@ -96,11 +97,20 @@ class Model:
         return data
 
     @staticmethod
-    def process_mask(mask: torch.Tensor) -> torch.Tensor:
+    def process_mask(mask: np.ndarray) -> np.ndarray:
+        logger.debug(f'{mask.shape=}')
+        if mask.shape != (512, 256, 3):
+            return None
+        colors = np.unique(mask.reshape(-1, 3), axis=0)
+        colors = set(map(tuple, colors.tolist()))
+        logger.debug(f'{colors=}')
+        logger.debug(f'{colors - set(COLOR_LIST)=}')
+
         seg_map = np.full(mask.shape[:-1], -1)
         for index, color in enumerate(COLOR_LIST):
             seg_map[np.sum(mask == color, axis=2) == 3] = index
-        assert (seg_map != -1).all()
+        if not (seg_map != -1).all():
+            return None
         return seg_map
 
     @staticmethod
@@ -140,6 +150,8 @@ class Model:
         logger.debug(f'{sample_steps=}')
         mask = label_image.copy()
         seg_map = self.process_mask(mask)
+        if seg_map is None:
+            return
         self.model.segm = torch.from_numpy(seg_map).unsqueeze(0).unsqueeze(
             0).to(self.model.device)
         self.model.generate_quantized_segm()