timecheck hetamaps

supervisely-ecosystem · Apr 22, 2024 · f5b09af · f5b09af
1 parent c6e8b67
commit f5b09af
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/main.py b/src/main.py
@@ -202,7 +202,6 @@ def main_func(team: TeamInfo, project: ProjectInfo):
     u.sew_chunks_to_json(stats, project_fs_dir, updated_classes)
 
     u.calculate_and_save_heatmaps(
-        datasets,
         project_fs_dir,
         heatmaps,
         heatmaps_image_ids,

diff --git a/src/utils.py b/src/utils.py
@@ -522,8 +522,8 @@ def _update_heatmaps_sample(heatmaps_figure_ids, heatmaps_image_ids, figs: List[
             heatmaps_image_ids[fig.dataset_id].add(fig.entity_id)
 
 
+@sly.timeit
 def calculate_and_save_heatmaps(
-    datasets: List[DatasetInfo],
     project_fs_dir: str,
     heatmaps: dtools.ClassesHeatmaps,
     heatmaps_image_ids: Dict[int, Set[int]],
@@ -547,7 +547,10 @@ def calculate_and_save_heatmaps(
                     filtered = [x for x in figs if x.id in heatmaps_figure_ids[x.class_id]]
                     heatmaps.update2(image, filtered, skip_broken_geometry=True)
                     pbar.update(1)
+    t = sly.TinyTimer()
+    sly.logger.debug(f"start to_image: {t.get_sec()} secs.")
     heatmaps.to_image(f"{project_fs_dir}/{heatmaps.basename_stem}.png")
+    sly.logger.debug(f"Heatmaps: {t.get_sec()} secs.")
 
 
 @sly.timeit