From 33014fd5e7b7cdf8eaa9d9a1dd4327dd3608f367 Mon Sep 17 00:00:00 2001
From: imhuwq <imhuwq@gmail.com>
Date: Thu, 4 Jan 2024 10:19:59 +0800
Subject: [PATCH] bugfix(finish_batch_add_annotation): fix parameter error

---
 deepdataspace/model/image.py               |  2 +-
 deepdataspace/plugins/coco2017/importer.py |  5 ++++-
 deepdataspace/utils/function.py            | 23 ++++++++++++++++++++++
 requirements-dev.txt                       |  1 +
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/deepdataspace/model/image.py b/deepdataspace/model/image.py
index 24e5466..2840232 100644
--- a/deepdataspace/model/image.py
+++ b/deepdataspace/model/image.py
@@ -461,7 +461,7 @@ def batch_add_annotation(self,
         self.objects.append(anno_obj)
 
     def finish_batch_add_annotation(self):
-        self.dataset.batch_save_image(self)
+        self.dataset.batch_save_image()
 
 
 _image_models: Dict[str, Type[ImageModel]] = {}  # a cache for ImageModel for each dataset
diff --git a/deepdataspace/plugins/coco2017/importer.py b/deepdataspace/plugins/coco2017/importer.py
index f50393e..ad727e6 100644
--- a/deepdataspace/plugins/coco2017/importer.py
+++ b/deepdataspace/plugins/coco2017/importer.py
@@ -33,6 +33,9 @@ def __init__(self, meta_path: str, enforce: bool = False):
 
         self.meta_path = os.path.abspath(meta_path)
         info = self.parse_meta(meta_path)
+        if info is None:
+            raise RuntimeError(f"Cannot import coco dataset: {meta_path}")
+
         dataset_name = info["dataset_name"]
         self.ground_truth = info["ground_truth"]
         self.image_root = info["image_root"]
@@ -148,7 +151,7 @@ def pre_run(self):
         self.format_data()
 
     def __iter__(self) -> Tuple[Dict, List[Dict]]:
-        for coco_image_data in self._images:
+        for coco_image_data in self._images[:1000000]:
             # data_sample = {"license"      : 2,
             #                "file_name"    : "000000000139.jpg",
             #                "coco_url"     : "http://images.cocodataset.org/val2017/000000000139.jpg",
diff --git a/deepdataspace/utils/function.py b/deepdataspace/utils/function.py
index ae5d9ad..8202c99 100644
--- a/deepdataspace/utils/function.py
+++ b/deepdataspace/utils/function.py
@@ -4,6 +4,8 @@
 Convenient functions about python function.
 """
 
+import cProfile
+import pstats
 import time
 from contextlib import contextmanager
 
@@ -32,6 +34,27 @@ def count_block_time(block_id: str, logger=print):
         logger(f"time cost of block[{block_id}]: {end - start}ms")
 
 
+@contextmanager
+def profile_perf(report_file: str, turn_on: bool = True):
+    """
+    Profile the performance of a code block, and save the result to report_file.
+    :param report_file: the target path where the performance data is saved to.
+    :param turn_on: only run the profile this is True or evaluates to True.
+    """
+    profile = cProfile.Profile()
+
+    if turn_on:
+        profile.enable()
+
+    try:
+        yield
+    finally:
+        if turn_on is True:
+            profile.disable()
+            stats = pstats.Stats(profile).sort_stats("cumulative")
+            stats.dump_stats(report_file)
+
+
 def retry(times: int, sleep: int = 0, exceptions: tuple = (Exception,)):
     """
     Retry a function or a method.
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 7c24060..c155d40 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -3,3 +3,4 @@ pytest==7.2.2
 pytest-cov==4.0.0
 Sphinx==5.3.0
 sphinx-rtd-theme==1.2.0
+snakeviz==2.2.0