Image copying in CVAT format (cvat-ai#2091)

minkoo01 · Sep 2, 2020 · ae6ec40 · ae6ec40
1 parent e52ff96
commit ae6ec40
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/cvat/apps/dataset_manager/formats/cvat.py b/cvat/apps/dataset_manager/formats/cvat.py
@@ -530,15 +530,16 @@ def _export(dst_file, task_data, anno_callback, save_images=False):
             frame_provider = FrameProvider(task_data.db_task.data)
             frames = frame_provider.get_frames(
                 frame_provider.Quality.ORIGINAL,
-                frame_provider.Type.NUMPY_ARRAY)
+                frame_provider.Type.BUFFER)
             for frame_id, (frame_data, _) in enumerate(frames):
                 frame_name = task_data.frame_info[frame_id]['path']
-                if '.' in frame_name:
-                    save_image(osp.join(img_dir, frame_name),
-                        frame_data, jpeg_quality=100, create_dir=True)
-                else:
-                    save_image(osp.join(img_dir, frame_name + '.png'),
-                        frame_data, create_dir=True)
+                ext = ''
+                if not '.' in osp.basename(frame_name):
+                    ext = '.png'
+                img_path = osp.join(img_dir, frame_name + ext)
+                os.makedirs(osp.dirname(img_path), exist_ok=True)
+                with open(img_path, 'wb') as f:
+                    f.write(frame_data.getvalue())
 
         make_zip_archive(temp_dir, dst_file)