Minor fixes

radekd91 · Jan 25, 2024 · efa04d0 · efa04d0
1 parent 43692d0
commit efa04d0
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/inferno/models/mica/MicaInputProcessing.py b/inferno/models/mica/MicaInputProcessing.py
@@ -193,7 +193,7 @@ def _dirty_image_preprocessing(self, input_image):
                 kps = kpss[bb_i]
 
             face = Face(bbox=bbox, kps=kps, det_score=det_score)
-            blob, aimg = get_arcface_input(face, img)
+            blob, aimg = get_arcface_input(face, img, image_is_bgr=False)
             aligned_image_list.append(aimg)
         aligned_images = np.array(aligned_image_list)
         # b,h,w,c to b,c,h,w

diff --git a/inferno/models/temporal/Preprocessors.py b/inferno/models/temporal/Preprocessors.py
@@ -261,6 +261,9 @@ def forward(self, batch, input_key, *args, output_prefix="gt_", test_time=False,
             for key in batch['landmarks'].keys():
                 batch_['landmarks'][key] = batch['landmarks'][key].view(B*T, -1, 2)
 
+            if 'mica_video' in batch: 
+                batch_['mica_video'] = batch['mica_video'].view(B*T, *batch['mica_video'].shape[-3:])
+
             values = self.model(batch_, training=False, validation=False)
         else:
             outputs = []
@@ -270,8 +273,21 @@ def forward(self, batch, input_key, *args, output_prefix="gt_", test_time=False,
                 batch_['landmarks'] = {}
                 for key in batch['landmarks'].keys():
                     batch_['landmarks'][key] = batch['landmarks'][key].view(B*T, -1, 2)[i:i+self.max_b]
+                if 'mica_video' in batch: 
+                    batch_['mica_video'] = batch['mica_video'].view(B*T, *batch['mica_video'].shape[-3:])[i:i+self.max_b]
                 out = self.model(batch_, training=False, validation=False)
                 outputs.append(out)
+
+                if 'image' in out: 
+                    del out['image']
+                if 'mica_images' in out: 
+                    del out['mica_images']
+                if 'predicted_image' in out: 
+                    del out['predicted_image']
+                if 'predicted_mask' in out: 
+                    del out['predicted_mask']
+                if 'albedo' in out: 
+                    del out['albedo']
 
             # combine into a single output
             values = cat_tensor_or_dict(outputs, dim=0)