improve handling of user input for not all sessions

src/vame/analysis/community_analysis.py

@@ -4,14 +4,16 @@
 import numpy as np
 from pathlib import Path
 import matplotlib.pyplot as plt
+from typing import List, Tuple
+
 from vame.util.auxiliary import read_config
 from vame.analysis.tree_hierarchy import (
     graph_to_tree,
     draw_tree,
     bag_nodes_by_cutline,
 )
 from vame.util.data_manipulation import consecutive
-from typing import List, Tuple
+from vame.util.cli import get_sessions_from_user_input
 from vame.schemas.states import save_state, CommunityFunctionSchema
 from vame.logging.logger import VameLogger
 from vame.schemas.project import SegmentationAlgorithms
@@ -558,29 +560,16 @@ def community(
         model_name = cfg["model_name"]
         n_clusters = cfg["n_clusters"]
 
-        sessions = []
-        if cfg["all_data"] == "No":
-            all_flag = input(
-                "Do you want to write motif videos for your entire dataset? \n"
-                "If you only want to use a specific dataset type filename: \n"
-                "yes/no/filename "
-            )
+        # Get sessions
+        if cfg["all_data"] in ["Yes", "yes"]:
+            sessions = cfg["session_names"]
         else:
-            all_flag = "yes"
-
-        if all_flag == "yes" or all_flag == "Yes":
-            for session in cfg["session_names"]:
-                sessions.append(session)
-        elif all_flag == "no" or all_flag == "No":
-            for session in cfg["session_names"]:
-                use_session = input("Do you want to quantify " + session + "? yes/no: ")
-                if use_session == "yes":
-                    sessions.append(session)
-                if use_session == "no":
-                    continue
-        # else:
-        #     files.append(all_flag)
+            sessions = get_sessions_from_user_input(
+                cfg=cfg,
+                action_message="run community analysis",
+            )
 
+        # Run community analysis for cohort=True
         if cohort:
             path_to_dir = Path(
                 os.path.join(

src/vame/analysis/generative_functions.py

@@ -6,9 +6,10 @@
 from sklearn.mixture import GaussianMixture
 
 from vame.schemas.states import GenerativeModelFunctionSchema, save_state
-from vame.util.auxiliary import read_config
 from vame.logging.logger import VameLogger
+from vame.util.auxiliary import read_config
 from vame.util.model_util import load_model
+from vame.util.cli import get_sessions_from_user_input
 from vame.schemas.project import SegmentationAlgorithms
 
 
@@ -257,29 +258,14 @@ def generative_model(
         model_name = cfg["model_name"]
         n_clusters = cfg["n_clusters"]
 
-        sessions = []
-        if cfg["all_data"] == "No":
-            all_flag = input(
-                "Do you want to write motif videos for your entire dataset? \n"
-                "If you only want to use a specific dataset type filename: \n"
-                "yes/no/filename "
-            )
+        # Get sessions
+        if cfg["all_data"] in ["Yes", "yes"]:
+            sessions = cfg["session_names"]
         else:
-            all_flag = "yes"
-
-        if all_flag == "yes" or all_flag == "Yes":
-            for session in cfg["session_names"]:
-                sessions.append(session)
-
-        elif all_flag == "no" or all_flag == "No":
-            for session in cfg["session_names"]:
-                use_session = input("Do you want to quantify " + session + "? yes/no: ")
-                if use_session == "yes":
-                    sessions.append(session)
-                if use_session == "no":
-                    continue
-        else:
-            sessions.append(all_flag)
+            sessions = get_sessions_from_user_input(
+                cfg=cfg,
+                action_message="generate samples",
+            )
 
         model = load_model(cfg, model_name, fixed=False)
 

src/vame/analysis/gif_creator.py

@@ -6,9 +6,11 @@
 import matplotlib
 from matplotlib import pyplot as plt
 from matplotlib.gridspec import GridSpec
+from typing import List, Tuple
+
 from vame.util.auxiliary import read_config
 from vame.util.gif_pose_helper import get_animal_frames
-from typing import List, Tuple
+from vame.util.cli import get_sessions_from_user_input
 from vame.logging.logger import VameLogger
 from vame.schemas.project import SegmentationAlgorithms
 
@@ -154,29 +156,14 @@ def gif(
     if segmentation_algorithm not in cfg["segmentation_algorithms"]:
         raise ValueError("Segmentation algorithm not found in config")
 
-    sessions = []
-    if cfg["all_data"] == "No":
-        all_flag = input(
-            "Do you want to write motif videos for your entire dataset? \n"
-            "If you only want to use a specific dataset type filename: \n"
-            "yes/no/filename "
-        )
-    else:
-        all_flag = "yes"
-
-    if all_flag == "yes" or all_flag == "Yes":
-        for session in cfg["session_names"]:
-            sessions.append(session)
-
-    elif all_flag == "no" or all_flag == "No":
-        for session in cfg["session_names"]:
-            use_session = input("Do you want to quantify " + session + "? yes/no: ")
-            if use_session == "yes":
-                sessions.append(session)
-            if use_session == "no":
-                continue
+    # Get sessions
+    if cfg["all_data"] in ["Yes", "yes"]:
+        sessions = cfg["session_names"]
     else:
-        sessions.append(all_flag)
+        sessions = get_sessions_from_user_input(
+            cfg=cfg,
+            action_message="create gifs",
+        )
 
     for session in sessions:
         path_to_file = os.path.join(

src/vame/analysis/pose_segmentation.py

@@ -5,13 +5,15 @@
 import numpy as np
 from pathlib import Path
 from typing import List, Tuple, Union
-from vame.util.data_manipulation import consecutive
 from hmmlearn import hmm
 from sklearn.cluster import KMeans
+
 from vame.schemas.states import save_state, SegmentSessionFunctionSchema
 from vame.logging.logger import VameLogger, TqdmToLogger
-from vame.util.auxiliary import read_config
 from vame.model.rnn_model import RNN_VAE
+from vame.util.auxiliary import read_config
+# from vame.util.data_manipulation import consecutive
+from vame.util.cli import get_sessions_from_user_input
 from vame.util.model_util import load_model
 
 
@@ -336,29 +338,14 @@ def segment_session(
                         )
                     )
 
-            sessions = []
-            if cfg["all_data"] == "No":
-                all_flag = input(
-                    "Do you want to quantify your entire dataset? \n"
-                    "If you only want to use a specific dataset type the session name: \n"
-                    "yes/no/session_name "
-                )
-            else:
-                all_flag = "yes"
-
-            if all_flag == "yes" or all_flag == "Yes":
+            # Get sessions
+            if cfg["all_data"] in ["Yes", "yes"]:
                 sessions = cfg["session_names"]
-            elif all_flag == "no" or all_flag == "No":
-                for session in cfg["session_names"]:
-                    use_session = input(
-                        "Do you want to quantify " + session + "? yes/no: "
-                    )
-                    if use_session == "yes":
-                        sessions.append(session)
-                    if use_session == "no":
-                        continue
             else:
-                sessions.append(all_flag)
+                sessions = get_sessions_from_user_input(
+                    cfg=cfg,
+                    action_message="run segmentation",
+                )
 
             use_gpu = torch.cuda.is_available()
             if use_gpu:

src/vame/analysis/umap.py

@@ -4,7 +4,9 @@
 from pathlib import Path
 import matplotlib.pyplot as plt
 from typing import Optional, Union
+
 from vame.util.auxiliary import read_config
+from vame.util.cli import get_sessions_from_user_input
 from vame.schemas.states import VisualizationFunctionSchema, save_state
 from vame.logging.logger import VameLogger
 from vame.schemas.project import SegmentationAlgorithms
@@ -264,27 +266,14 @@ def visualization(
         model_name = cfg["model_name"]
         n_clusters = cfg["n_clusters"]
 
-        sessions = []
-        if cfg["all_data"] == "No":
-            all_flag = input(
-                "Do you want to write motif videos for your entire dataset? \n"
-                "If you only want to use a specific dataset, type the session name: \n"
-                "yes/no/session_name "
-            )
-        else:
-            all_flag = "yes"
-
-        if all_flag == "yes" or all_flag == "Yes":
+        # Get sessions
+        if cfg["all_data"] in ["Yes", "yes"]:
             sessions = cfg["session_names"]
-        elif all_flag == "no" or all_flag == "No":
-            for session in cfg["session_names"]:
-                use_session = input("Do you want to quantify " + session + "? yes/no: ")
-                if use_session == "yes":
-                    sessions.append(session)
-                if use_session == "no":
-                    continue
         else:
-            sessions.append(all_flag)
+            sessions = get_sessions_from_user_input(
+                cfg=cfg,
+                action_message="generate visualization",
+            )
 
         for idx, session in enumerate(sessions):
             path_to_file = os.path.join(

src/vame/analysis/videowriter.py

@@ -4,13 +4,15 @@
 import cv2 as cv
 import tqdm
 from typing import Union
+import imageio
+
 from vame.util.auxiliary import read_config
+from vame.util.cli import get_sessions_from_user_input
 from vame.schemas.states import (
     save_state,
     MotifVideosFunctionSchema,
     CommunityVideosFunctionSchema,
 )
-import imageio
 from vame.logging.logger import VameLogger, TqdmToLogger
 from vame.schemas.project import SegmentationAlgorithms
 
@@ -214,29 +216,15 @@ def motif_videos(
         n_clusters = cfg["n_clusters"]
 
         logger.info(f"Creating motif videos for algorithm: {segmentation_algorithm}...")
-        sessions = []
-        if cfg["all_data"] == "No":
-            all_flag = input(
-                "Do you want to write motif videos for your entire dataset? \n"
-                "If you only want to use a specific dataset, type the session name: \n"
-                "yes/no/session_name"
-            )
-        else:
-            all_flag = "yes"
-
-        if all_flag == "yes" or all_flag == "Yes":
-            for session in cfg["session_names"]:
-                sessions.append(session)
-
-        elif all_flag == "no" or all_flag == "No":
-            for session in cfg["session_names"]:
-                use_session = input("Do you want to quantify " + session + "? yes/no: ")
-                if use_session == "yes":
-                    sessions.append(session)
-                if use_session == "no":
-                    continue
+
+        # Get sessions
+        if cfg["all_data"] in ["Yes", "yes"]:
+            sessions = cfg["session_names"]
         else:
-            sessions.append(all_flag)
+            sessions = get_sessions_from_user_input(
+                cfg=cfg,
+                action_message="write motif videos",
+            )
 
         logger.info("Cluster size is: %d " % n_clusters)
         for session in sessions:
@@ -329,27 +317,14 @@ def community_videos(
         model_name = cfg["model_name"]
         n_clusters = cfg["n_clusters"]
 
-        sessions = []
-        if cfg["all_data"] == "No":
-            all_flag = input(
-                "Do you want to write motif videos for your entire dataset? \n"
-                "If you only want to use a specific dataset, type the session name: \n"
-                "yes/no/session_name"
-            )
-        else:
-            all_flag = "yes"
-
-        if all_flag == "yes" or all_flag == "Yes":
+        # Get sessions
+        if cfg["all_data"] in ["Yes", "yes"]:
             sessions = cfg["session_names"]
-        elif all_flag == "no" or all_flag == "No":
-            for session in cfg["session_names"]:
-                use_session = input("Do you want to quantify " + session + "? yes/no: ")
-                if use_session == "yes":
-                    sessions.append(session)
-                if use_session == "no":
-                    continue
         else:
-            sessions.append(all_flag)
+            sessions = get_sessions_from_user_input(
+                cfg=cfg,
+                action_message="write community videos",
+            )
 
         logger.info("Cluster size is: %d " % n_clusters)
         for session in sessions:

src/vame/util/cli.py

@@ -0,0 +1,29 @@
+from typing import List
+
+
+def get_sessions_from_user_input(
+    cfg: dict,
+    action_message: str = "run this step",
+) -> List[str]:
+    user_input = input(
+        f"Do you want to {action_message} for your entire dataset? \n"
+        "If you only want to use a specific session, type the session name \n"
+        "yes/no/<session_name>: "
+    )
+    if user_input in ["Yes", "yes"]:
+        sessions = cfg["session_names"]
+    elif user_input in ["No", "no"]:
+        sessions = []
+        for session in cfg["session_names"]:
+            use_session = input("Do you want to use " + session + "? yes/no: ")
+            if use_session in ["Yes", "yes"]:
+                sessions.append(session)
+            else:
+                continue
+    else:
+        if user_input in cfg["session_names"]:
+            sessions = [user_input]
+        else:
+            raise ValueError(
+                "Invalid input. Please enter yes, no, or a valid session name."
+            )