Upgrade some union, list, dict annotations to 3.10 (#239)

* Upgrade some union, list, dict annotations to 3.10

* specify 3.10 as minimum python version

* fixup setup.py

* fixup predictions.py

* remove more list, union

* remove int

Author: ksikka

cli/friendly.py

@@ -1,7 +1,6 @@
 import argparse
 import shutil
 import sys
-from typing import List
 
 
 class ArgumentParser(argparse.ArgumentParser):
@@ -39,14 +38,14 @@ def __init__(self, *args, **kwargs):
 class _HelpFormatter(argparse.HelpFormatter):
     """Modifications on help text formatting for easier readability."""
 
-    def _split_lines(self, text: str, width: int) -> List[str]:
+    def _split_lines(self, text: str, width: int) -> list[str]:
         """Modified to preserve newlines and long words."""
         # First split into paragraphs, then wrap each separately:
         # https://docs.python.org/3/library/textwrap.html#textwrap.TextWrapper.replace_whitespace
         paragraphs = text.splitlines()
         import textwrap
 
-        lines: List[str] = []
+        lines: list[str] = []
         for p in paragraphs:
             p_lines = textwrap.wrap(
                 p, width, break_long_words=False, break_on_hyphens=False

lightning_pose/__init__.py

@@ -1,6 +1,7 @@
-__version__ = "1.6.1"
+__version__ = "1.7.0"
 
 from pathlib import Path
+
 from omegaconf import OmegaConf
 
 LP_ROOT_PATH = (Path(__file__).parent.parent).absolute()

lightning_pose/apps/utils.py

@@ -3,7 +3,7 @@
 import os
 from collections import defaultdict
 from pathlib import Path
-from typing import Dict, List, Tuple
+from typing import Tuple
 
 import pandas as pd
 import streamlit as st
@@ -16,7 +16,7 @@
 
 
 @st.cache_resource
-def update_labeled_file_list(model_preds_folders: List[str], use_ood: bool = False) -> List[list]:
+def update_labeled_file_list(model_preds_folders: list[str], use_ood: bool = False) -> list[list]:
     per_model_preds = []
     for model_pred_folder in model_preds_folders:
         # pull labeled results from each model folder
@@ -42,9 +42,9 @@ def update_labeled_file_list(model_preds_folders: List[str], use_ood: bool = Fal
 @st.cache_resource
 def update_vid_metric_files_list(
     video: str,
-    model_preds_folders: List[str],
+    model_preds_folders: list[str],
     video_subdir: str = "video_preds",
-) -> List[list]:
+) -> list[list]:
     per_vid_preds = []
     for model_preds_folder in model_preds_folders:
         # pull each prediction file associated with a particular video
@@ -65,7 +65,7 @@ def update_vid_metric_files_list(
 
 
 @st.cache_resource
-def get_all_videos(model_preds_folders: List[str], video_subdir: str = "video_preds") -> list:
+def get_all_videos(model_preds_folders: list[str], video_subdir: str = "video_preds") -> list:
     # find each video that is predicted on by the models
     # wrap in Path so that it looks like an UploadedFile object
     # returned by streamlit's file_uploader
@@ -89,7 +89,7 @@ def get_all_videos(model_preds_folders: List[str], video_subdir: str = "video_pr
 
 
 @st.cache_data
-def concat_dfs(dframes: Dict[str, pd.DataFrame]) -> Tuple[pd.DataFrame, List[str]]:
+def concat_dfs(dframes: dict[str, pd.DataFrame]) -> Tuple[pd.DataFrame, list[str]]:
     counter = 0
     for model_name, dframe in dframes.items():
         mask = dframe.columns.get_level_values("coords").isin(["x", "y", "likelihood"])
@@ -143,7 +143,7 @@ def get_df_scatter(
     return pd.concat(df_scatters)
 
 
-def get_col_names(keypoint: str, coordinate: str, models: List[str]) -> List[str]:
+def get_col_names(keypoint: str, coordinate: str, models: list[str]) -> list[str]:
     return [get_full_name(keypoint, coordinate, model) for model in models]
 
 
@@ -162,7 +162,7 @@ def get_full_name(keypoint: str, coordinate: str, model: str) -> str:
 # ----------------------------------------------
 @st.cache_data
 def build_precomputed_metrics_df(
-    dframes: Dict[str, pd.DataFrame], keypoint_names: List[str], **kwargs,
+    dframes: dict[str, pd.DataFrame], keypoint_names: list[str], **kwargs,
 ) -> dict:
     concat_dfs = defaultdict(list)
     for model_name, df_dict in dframes.items():
@@ -194,7 +194,7 @@ def build_precomputed_metrics_df(
 
 @st.cache_data
 def get_precomputed_error(
-    df: pd.DataFrame, keypoint_names: List[str], model_name: str,
+    df: pd.DataFrame, keypoint_names: list[str], model_name: str,
 ) -> pd.DataFrame:
     # collect results
     df_ = df
@@ -206,7 +206,7 @@ def get_precomputed_error(
 
 @st.cache_data
 def compute_confidence(
-    df: pd.DataFrame, keypoint_names: List[str], model_name: str, **kwargs,
+    df: pd.DataFrame, keypoint_names: list[str], model_name: str, **kwargs,
 ) -> pd.DataFrame:
 
     if df.shape[1] % 3 == 1:
@@ -236,7 +236,7 @@ def get_model_folders(
     model_dir: str,
     require_predictions: bool = True,
     require_tb_logs: bool = False,
-) -> List[str]:
+) -> list[str]:
     """Find all model folders in a higher-level directory, conditional on directory contents."""
     # strip trailing slash if present
     if model_dir[-1] == os.sep:
@@ -260,7 +260,7 @@ def get_model_folders(
 
 
 # just to get the last two levels of the path
-def get_model_folders_vis(model_folders: List[str]) -> List[str]:
+def get_model_folders_vis(model_folders: list[str]) -> list[str]:
     fs = []
     for f in model_folders:
         fs.append(f.split("/")[-2:])

lightning_pose/data/dali.py

@@ -1,7 +1,7 @@
 """Data pipelines based on efficient video reading by nvidia dali package."""
 
 import os
-from typing import Dict, List, Literal, Optional, Union
+from typing import Literal
 
 import numpy as np
 import nvidia.dali.fn as fn
@@ -25,14 +25,14 @@
 # cannot typecheck due to way pipeline_def decorator consumes additional args
 @pipeline_def
 def video_pipe(
-    filenames: Union[List[str], str],
-    resize_dims: Optional[List[int]] = None,
+    filenames: list[str] | str,
+    resize_dims: list[int] | None = None,
     random_shuffle: bool = False,
     sequence_length: int = 16,
     pad_sequences: bool = True,
     initial_fill: int = 16,
-    normalization_mean: List[float] = _IMAGENET_MEAN,
-    normalization_std: List[float] = _IMAGENET_STD,
+    normalization_mean: list[float] = _IMAGENET_MEAN,
+    normalization_std: list[float] = _IMAGENET_STD,
     name: str = "reader",
     step: int = 1,
     pad_last_batch: bool = False,
@@ -174,7 +174,7 @@ def __len__(self) -> int:
     @staticmethod
     def _dali_output_to_tensors(
         batch: list
-    ) -> Union[UnlabeledBatchDict, MultiviewUnlabeledBatchDict]:
+    ) -> UnlabeledBatchDict | MultiviewUnlabeledBatchDict:
 
         # always batch_size=1
 
@@ -229,7 +229,7 @@ def _dali_output_to_tensors(
                 frames=frames, transforms=transforms, bbox=bbox, is_multiview=True,
             )
 
-    def __next__(self) -> Union[UnlabeledBatchDict, MultiviewUnlabeledBatchDict]:
+    def __next__(self) -> UnlabeledBatchDict | MultiviewUnlabeledBatchDict:
         batch = super().__next__()
         return self._dali_output_to_tensors(batch=batch)
 
@@ -245,10 +245,10 @@ def __init__(
         self,
         train_stage: Literal["predict", "train"],
         model_type: Literal["base", "context"],
-        filenames: Union[List[str], List[List[str]]],
-        resize_dims: List[int],
-        dali_config: Union[dict, DictConfig] = None,
-        imgaug: Optional[str] = "default",
+        filenames: list[str] | list[list[str]],
+        resize_dims: list[int],
+        dali_config: dict | DictConfig = None,
+        imgaug: str | None = "default",
         num_threads: int = 1,
     ) -> None:
 
@@ -315,9 +315,9 @@ def num_iters(self) -> int:
 
     def _setup_pipe_dict(
         self,
-        filenames: Union[List[str], List[List[str]]],
+        filenames: list[str] | list[list[str]],
         imgaug: str,
-    ) -> Dict[str, dict]:
+    ) -> dict[str, dict]:
         """All of the pipeline args in one place."""
         # When running with multiple GPUs, the LOCAL_RANK variable correctly
         # contains the DDP Local Rank, which is also the cuda device index.

lightning_pose/data/datamodules.py

@@ -1,7 +1,7 @@
 """Data modules split a dataset into train, val, and test modules."""
 
 import copy
-from typing import List, Literal, Optional, Union
+from typing import Literal
 
 import imgaug.augmenters as iaa
 import lightning.pytorch as pl
@@ -36,9 +36,9 @@ def __init__(
         test_batch_size: int = 1,
         num_workers: int = 8,
         train_probability: float = 0.8,
-        val_probability: Optional[float] = None,
-        test_probability: Optional[float] = None,
-        train_frames: Optional[Union[float, int]] = None,
+        val_probability: float | None = None,
+        test_probability: float | None = None,
+        train_frames: float | int | None = None,
         torch_seed: int = 42,
     ) -> None:
         """Data module splits a dataset into train, val, and test data loaders.
@@ -75,7 +75,7 @@ def __init__(
         self.test_dataset = None  # populated by self.setup()
         self.torch_seed = torch_seed
 
-    def setup(self, stage: Optional[str] = None) -> None:  # stage arg needed for ptl
+    def setup(self, stage: str | None = None) -> None:  # stage arg needed for ptl
 
         datalen = self.dataset.__len__()
         print(f"Number of labeled images in the full dataset (train+val+test): {datalen}")
@@ -170,17 +170,17 @@ class UnlabeledDataModule(BaseDataModule):
     def __init__(
         self,
         dataset: torch.utils.data.Dataset,
-        video_paths_list: Union[List[str], str],
-        dali_config: Union[dict, DictConfig],
-        view_names: Optional[List[str]] = None,
+        video_paths_list: list[str] | str,
+        dali_config: dict | DictConfig,
+        view_names: list[str] | None = None,
         train_batch_size: int = 16,
         val_batch_size: int = 16,
         test_batch_size: int = 1,
         num_workers: int = 8,
         train_probability: float = 0.8,
-        val_probability: Optional[float] = None,
-        test_probability: Optional[float] = None,
-        train_frames: Optional[float] = None,
+        val_probability: float | None = None,
+        test_probability: float | None = None,
+        train_frames: float | None = None,
         torch_seed: int = 42,
         imgaug: Literal["default", "dlc", "dlc-top-down"] = "default",
     ) -> None:

lightning_pose/data/datasets.py

@@ -3,7 +3,7 @@
 
 import os
 from pathlib import Path
-from typing import Callable, List, Literal, Optional, Tuple, Union
+from typing import Callable, List, Literal, Tuple
 
 import numpy as np
 import pandas as pd
@@ -36,8 +36,8 @@ def __init__(
         self,
         root_directory: str,
         csv_path: str,
-        header_rows: Optional[List[int]] = [0, 1, 2],
-        imgaug_transform: Optional[Callable] = None,
+        header_rows: list[int] | None = [0, 1, 2],
+        imgaug_transform: Callable | None = None,
         do_context: bool = False,
     ) -> None:
         """Initialize a dataset for regression (rather than heatmap) models.
@@ -193,8 +193,8 @@ def __init__(
         self,
         root_directory: str,
         csv_path: str,
-        header_rows: Optional[List[int]] = [0, 1, 2],
-        imgaug_transform: Optional[Callable] = None,
+        header_rows: list[int] | None = [0, 1, 2],
+        imgaug_transform: Callable | None = None,
         downsample_factor: Literal[1, 2, 3] = 2,
         do_context: bool = False,
         uniform_heatmaps: bool = False,
@@ -308,13 +308,13 @@ class MultiviewHeatmapDataset(torch.utils.data.Dataset):
     def __init__(
         self,
         root_directory: str,
-        csv_paths: List[str],
-        view_names: List[str],
-        header_rows: Optional[List[int]] = [0, 1, 2],
+        csv_paths: list[str],
+        view_names: list[str],
+        header_rows: list[int] | None = [0, 1, 2],
         downsample_factor: Literal[1, 2, 3] = 2,
         uniform_heatmaps: bool = False,
         do_context: bool = False,
-        imgaug_transform: Optional[Callable] = None,
+        imgaug_transform: Callable | None = None,
     ) -> None:
         """Initialize the MultiViewHeatmap Dataset.
 
@@ -424,10 +424,12 @@ def output_shape(self) -> tuple:
     def num_views(self) -> int:
         return len(self.view_names)
 
-    def fusion(self, datadict: dict) -> Tuple[
-        Union[
-            TensorType["num_views", "RGB":3, "image_height", "image_width", float],
-            TensorType["num_views", "frames", "RGB":3, "image_height", "image_width", float]
+    def fusion(
+        self, datadict: dict
+    ) -> Tuple[
+        TensorType["num_views", "RGB":3, "image_height", "image_width", float]
+        | TensorType[
+            "num_views", "frames", "RGB":3, "image_height", "image_width", float
         ],
         TensorType["keypoints"],
         TensorType["num_views", "heatmap_height", "heatmap_width", float],
@@ -489,6 +491,6 @@ def __getitem__(self, idx: int) -> MultiviewHeatmapLabeledExampleDict:
             bbox=bboxes,
             idxs=idx,
             num_views=self.num_views,  # int
-            concat_order=concat_order,  # List[str]
-            view_names=self.view_names,  # List[str]
+            concat_order=concat_order,  # list[str]
+            view_names=self.view_names,  # list[str]
         )