Fix loading file and image processing in CIFAR (#284)

* Fix image layout and encoding problems

* Update Changelog

Co-authored-by: Maxim Zhiltsov <maxim.zhiltsov@intel.com>

CHANGELOG.md

@@ -22,7 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 -
 
 ### Fixed
--
+- Incorrect image layout on saving and a problem with ecoding on loading (<https://github.com/openvinotoolkit/datumaro/pull/284>)
 
 ### Security
 -

datumaro/plugins/cifar_format.py

@@ -69,7 +69,7 @@ def _load_items(self, path):
         # 'filenames': list
         # 'labels': list
         with open(path, 'rb') as anno_file:
-            annotation_dict = pickle.load(anno_file)
+            annotation_dict = pickle.load(anno_file, encoding='latin1')
 
         labels = annotation_dict.get('labels', [])
         filenames = annotation_dict.get('filenames', [])
@@ -94,11 +94,13 @@ def _load_items(self, path):
             if 0 < len(images_data):
                 image = images_data[i]
                 if size is not None and image is not None:
-                    image = image.reshape(size[i][0],
-                        size[i][1], 3).astype(np.uint8)
+                    image = image.reshape(3, size[i][0],
+                        size[i][1]).astype(np.uint8)
+                    image = np.transpose(image, (1, 2, 0))
                 elif image is not None:
-                    image = image.reshape(CifarPath.IMAGE_SIZE,
-                        CifarPath.IMAGE_SIZE, 3).astype(np.uint8)
+                    image = image.reshape(3, CifarPath.IMAGE_SIZE,
+                        CifarPath.IMAGE_SIZE).astype(np.uint8)
+                    image = np.transpose(image, (1, 2, 0))
 
             items[item_id] = DatasetItem(id=item_id, subset=self._subset,
                 image=image, annotations=annotations)
@@ -150,7 +152,8 @@ def apply(self):
                         data.append(None)
                     else:
                         image = image.data
-                        data.append(image.reshape(-1).astype(np.uint8))
+                        data.append(np.transpose(image,
+                            (2, 0, 1)).reshape(-1).astype(np.uint8))
                         if image.shape[0] != CifarPath.IMAGE_SIZE or \
                                 image.shape[1] != CifarPath.IMAGE_SIZE:
                             image_sizes[len(data) - 1] = (image.shape[0], image.shape[1])

tests/test_cifar_format.py

@@ -146,12 +146,17 @@ def test_can_import(self):
             DatasetItem(id='image_4', subset='test',
                 image=np.ones((32, 32, 3)),
                 annotations=[Label(2)]
+            ),
+            DatasetItem(id='image_5', subset='test',
+                image=np.array([[[1., 2., 3.], [4., 5., 6.]],
+                                [[1., 2., 3.], [4., 5., 6.]]]),
+                annotations=[Label(3)]
             )
         ], categories=['airplane', 'automobile', 'bird', 'cat'])
 
         dataset = Dataset.import_from(DUMMY_DATASET_DIR, 'cifar')
 
-        compare_datasets(self, expected_dataset, dataset)
+        compare_datasets(self, expected_dataset, dataset, require_images=True)
 
     @mark_requirement(Requirements.DATUM_GENERAL_REQ)
     def test_can_detect(self):