add preview mode handler

supervisely-ecosystem · Sep 16, 2024 · f289a49 · f289a49
1 parent 3de559a
commit f289a49
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 14 deletions.
diff --git a/src/compute/layers/processing/SplitDataLayer.py b/src/compute/layers/processing/SplitDataLayer.py
@@ -109,16 +109,19 @@ def _split_by_tags() -> List[Tuple[ImageDescriptor, Annotation]]:
                     items.append((new_label_item_desc, ann))
             return items
 
-        item_desc, ann = data_el
-        item_idx = item_desc.get_item_idx()
-
-        split_func_map = {
-            "percent": _split_by_percent,
-            "number": _split_by_num,
-            "classes": _split_by_class,
-            "tags": _split_by_tags,
-        }
-        split_method = self.settings["split_method"]
-        func = split_func_map.get(split_method)
-        items = func()
-        yield from items
+        if self.net.preview_mode:
+            yield data_el
+        else:
+            item_desc, ann = data_el
+            item_idx = item_desc.get_item_idx()
+
+            split_func_map = {
+                "percent": _split_by_percent,
+                "number": _split_by_num,
+                "classes": _split_by_class,
+                "tags": _split_by_tags,
+            }
+            split_method = self.settings["split_method"]
+            func = split_func_map.get(split_method)
+            items = func()
+            yield from items
diff --git a/src/compute/main.py b/src/compute/main.py
@@ -122,7 +122,7 @@ def main(
         raise e
 
     total = net.get_total_elements()
-    net.total_element_cnt = total
+    net.total_elements_cnt = total
     if total == 0:
         raise GraphError(
             "There are no elements to process. Make sure that you selected input project and it's not empty"