allow avi extension

YttriLab · Nov 9, 2023 · 41c258c · 41c258c
1 parent 0955577
commit 41c258c
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 10 deletions.
diff --git a/asoid/apps/D_manual_active_learning.py b/asoid/apps/D_manual_active_learning.py
@@ -520,15 +520,24 @@ def prompt_setup(software, ftype, selected_bodyparts, annotation_classes,
                                                , help = PLAYBACK_SPEED_HELP)
 
             col1_exp.write(f'Equivalent to {round(output_fps / framerate, 2)} X speed')
-            frame_dir = os.path.join(videos_dir,
-                                     str.join('', (st.session_state['uploaded_vid'].name.rpartition('.mp4')[0],
-                                                   '_pngs')))
+            if st.session_state['uploaded_vid'].name.endswith('mp4'):
+                frame_dir = os.path.join(videos_dir,
+                                         str.join('', (st.session_state['uploaded_vid'].name.rpartition('.mp4')[0],
+                                                       '_pngs')))
+            else:
+                frame_dir = os.path.join(videos_dir,
+                                         str.join('', (st.session_state['uploaded_vid'].name.rpartition('.avi')[0],
+                                                       '_pngs')))
             os.makedirs(frame_dir, exist_ok=True)
             col3_exp.success(f'Entered **{frame_dir}** as the frame directory.')
-
-            shortvid_dir = os.path.join(project_dir, iter_dir,
-                                        str.join('', (st.session_state['uploaded_vid'].name.rpartition('.mp4')[0],
-                                                      '_refine_vids')))
+            if st.session_state['uploaded_vid'].name.endswith('mp4'):
+                shortvid_dir = os.path.join(project_dir, iter_dir,
+                                            str.join('', (st.session_state['uploaded_vid'].name.rpartition('.mp4')[0],
+                                                          '_refine_vids')))
+            else:
+                shortvid_dir = os.path.join(project_dir, iter_dir,
+                                            str.join('', (st.session_state['uploaded_vid'].name.rpartition('.avi')[0],
+                                                          '_refine_vids')))
             os.makedirs(shortvid_dir, exist_ok=True)
             col3_exp.success(f'Entered **{shortvid_dir}** as the refined video directory.')
 
@@ -631,7 +640,12 @@ def main(ri=None, config=None):
         refined_vid_dirs.extend(['Add New Video'])
         outlier_methods = ['Random', 'Low Confidence']
         try:
-            new_vid_name = str.join('', (st.session_state['uploaded_vid'].name.rpartition('.mp4')[0], '_refine_vids'))
+            if st.session_state['uploaded_vid'].name.endswith('mp4'):
+                new_vid_name = str.join('',
+                                        (st.session_state['uploaded_vid'].name.rpartition('.mp4')[0], '_refine_vids'))
+            else:
+                new_vid_name = str.join('',
+                                        (st.session_state['uploaded_vid'].name.rpartition('.avi')[0], '_refine_vids'))
             selected_refine_dir = ri.radio('Select Refinement', refined_vid_dirs,
                                            index=refined_vid_dirs.index(new_vid_name),
                                            horizontal=True, key='selected_refine'

diff --git a/asoid/apps/E_create_new_training.py b/asoid/apps/E_create_new_training.py
@@ -10,18 +10,18 @@
 
 def create_new_training_features_targets(project_dir, selected_iter, new_features, new_targets):
     # load existing training
-    # st.write(new_features.shape, new_targets.shape)
+    # st.write(project_dir, selected_iter)
     iter_folder = str.join('', ('iteration-', str(selected_iter)))
     [features, targets, frames2integ] = \
         load_features(project_dir, iter_folder)
+    # st.write(np.unique(targets, return_counts=True))
     # st.write(features.shape, targets.shape)
     # add iteration number
     new_iter_folder = str.join('', ('iteration-', str(selected_iter + 1)))
     os.makedirs(os.path.join(project_dir, new_iter_folder), exist_ok=True)
     # incorporate new features/targets into existing training
     appended_features = np.vstack((features, new_features))
     appended_targets = np.hstack((targets, new_targets))
-    # st.write(appended_features.shape, appended_targets.shape)
     # save into new iteration folder
     save_data(project_dir, new_iter_folder, 'feats_targets.sav',
               [
@@ -107,6 +107,7 @@ def main(ri=None, config=None):
                 new_targets_dir.append(np.hstack(new_targets_byclass))
             new_features = np.vstack(new_features_dir)
             new_targets = np.hstack(new_targets_dir)
+            # st.write(np.unique(new_targets, return_counts=True))
         except:
             pass