From 5f94d104969636ca34cb52389c1c3cd6d9aa3e57 Mon Sep 17 00:00:00 2001
From: Thomas Varsavsky <tom@v7labs.com>
Date: Fri, 15 Dec 2023 15:05:46 +0000
Subject: [PATCH] added type hints

---
 darwin/importer/formats/nifti.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/darwin/importer/formats/nifti.py b/darwin/importer/formats/nifti.py
index 671b5f397..018690188 100644
--- a/darwin/importer/formats/nifti.py
+++ b/darwin/importer/formats/nifti.py
@@ -196,7 +196,9 @@ def get_polygon_video_annotations(
         raise Exception("If is_mpr is True, slot_names must be of length 3")
 
 
-def get_mask_video_annotations(volume, processed_class_map, slot_names):
+def get_mask_video_annotations(
+    volume: np.ndarray, processed_class_map: Dict, slot_names: List[str]
+) -> Optional[List[dt.VideoAnnotation]]:
     """
     The function takes a volume and a class map and returns a list of video annotations
 
@@ -280,13 +282,13 @@ def get_mask_video_annotations(volume, processed_class_map, slot_names):
 
 
 def nifti_to_video_polygon_annotation(
-    volume,
-    class_name,
-    class_idxs,
-    slot_names,
-    view_idx=2,
-    pixdims=(1, 1, 1),
-):
+    volume: np.ndarray,
+    class_name: str,
+    class_idxs: List[int],
+    slot_names: List[str],
+    view_idx: int = 2,
+    pixdims: Tuple[int, int, int] = (1, 1, 1),
+) -> Optional[List[dt.VideoAnnotation]]:
     frame_annotations = OrderedDict()
     for i in range(volume.shape[view_idx]):
         if view_idx == 2:
@@ -485,7 +487,7 @@ def process_nifti(
         return data_array, pixdims
 
 
-def convert_to_dense_rle(raster):
+def convert_to_dense_rle(raster: np.ndarray) -> List[int]:
     dense_rle, prev_val, cnt = [], None, 0
     for val in raster.T.flat:
         if val == prev_val: