Convert Shapes to lists, add extra dimension to images (#3455)

Author: anzhella-pankratova

demos/image_translation_demo/python/image_translation_demo/models.py

@@ -21,8 +21,8 @@ def __init__(self, core, model_path, device='CPU'):
 
         inputs = [node.get_any_name() for node in model.inputs]
         self.input_semantics, self.reference_image, self.reference_semantics = inputs
-        self.input_semantic_size = model.input(self.input_semantics).shape
-        self.input_image_size = model.input(self.reference_image).shape
+        self.input_semantic_size = list(model.input(self.input_semantics).shape)
+        self.input_image_size = list(model.input(self.reference_image).shape)
 
         compiled_model = core.compile_model(model, device)
         self.output_tensor = compiled_model.outputs[0]
@@ -46,7 +46,7 @@ def __init__(self, core, model_path, device='CPU'):
             raise RuntimeError("The SegmentationModel expects 1 output layer")
 
         self.input_tensor_name = model.inputs[0].get_any_name()
-        self.input_size = model.inputs[0].shape
+        self.input_size = list(model.inputs[0].shape)
 
         compiled_model = core.compile_model(model, device)
         self.output_tensor = compiled_model.outputs[0]

demos/image_translation_demo/python/image_translation_demo/preprocessing.py

@@ -34,16 +34,19 @@ def preprocess_semantics(semantic_mask, input_size):
                                interpolation=cv2.INTER_NEAREST)
     # create one-hot label map
     semantic_mask = scatter(semantic_mask, classes=input_size[1])
+
+    if len(semantic_mask.shape) == 3:
+        return np.expand_dims(semantic_mask, axis=0)
     return semantic_mask
 
 
 def preprocess_for_seg_model(image, input_size):
     image = cv2.resize(image, dsize=tuple(input_size[2:]), interpolation=cv2.INTER_LINEAR)
     image = np.transpose(image, (2, 0, 1))
-    return image
+    return np.expand_dims(image, axis=0)
 
 
 def preprocess_image(image, input_size):
     image = cv2.resize(image, dsize=tuple(input_size[2:]), interpolation=cv2.INTER_CUBIC)
     image = np.transpose(image, (2, 0, 1))
-    return image
+    return np.expand_dims(image, axis=0)