diff --git a/inferno/models/mica/MicaInputProcessing.py b/inferno/models/mica/MicaInputProcessing.py index dd67202b..85c56188 100644 --- 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 index 6c0de77e..e2279b8a 100644 --- 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)