WIP: Making the demo work on hf spaces

app.py

@@ -1,7 +1,8 @@
 """Streamlit demo to visualize auto-annotated Foley segments from movie clips."""
 import os
-from os.path import join, exists, dirname, abspath
+from os.path import join, exists, dirname, abspath, basename
 import json
+from glob import glob
 
 from tqdm import tqdm
 import numpy as np
@@ -147,78 +148,72 @@ if __name__ == "__main__":
         "**Instructions**: Click the **Reload** button to see segments from a new clip. "\
         "Reloading the page is not necessary."
     )
-    
-    use_local = False
-    data_root = get_data_root_from_hostname()
-    data_dir = join(data_root, "CondensedMovies")
-    video_dir = join(data_dir, "pytube_videos")
-
-    annot_dir = join(repo_path, "external/CondensedMovies/data/metadata/")
-
-    if "subdf" not in st.session_state:
-        df = load_clips_df(join(".", "clips.csv"), data_dir, verbose=True)
-        df["annot_filtered"] = df["annot_path"].apply(lambda x: x.replace(".json", "_filtered.json"))
-        df = df[df["annot_filtered"].apply(exists)]
-        df["num_foley_segments"] = df["annot_filtered"].apply(lambda f: sum(load_json(f)["keep_status"]))
-        subdf = df[df["num_foley_segments"].apply(lambda x: x > 0)]
-        st.session_state.subdf = subdf
-        num_foley = subdf["num_foley_segments"].sum()
-        st.session_state.num_foley = num_foley
-        print("Loaded subdf with {} rows".format(len(subdf)))
+
+    csv_path = "./clips.csv"
+    ann_dirs = glob(join(".", "annotations_", "*"))
+    annot_paths = glob(join(".", "annotations_*", "*_filtered.json"))
+    print("Total number of clips: {}".format(len(annot_paths)))
+
+
+    if "data" not in st.session_state:
+        # store video ids
+        video_ids = [basename(x).split("_filtered.json")[0] for x in annot_paths]
+        
+        # load annotation data
+        data = [load_json(p) for p in annot_paths]
+        num_foley_per_clip = [sum(d["keep_status"]) for d in data]
+        num_foley_segments = np.sum(num_foley_per_clip)
+        data = [d for d, n in zip(data, num_foley_per_clip) if n > 0]
+        
+        # store variables
+        st.session_state.video_ids = video_ids
+        st.session_state.data = data
+        st.session_state.num_foley_segments = num_foley_segments
 
     
     reload_button = st.button("Reload")
-    # index = 0
-    index = np.random.randint(0, len(st.session_state.subdf))
+    index = np.random.randint(0, len(st.session_state.data))
     if reload_button:
-        index = np.random.randint(0, len(st.session_state.subdf))
+        index = np.random.randint(0, len(st.session_state.data))
+
+    # Gather data
+    annot = st.session_state.data[index]
+    video_id = st.session_state.video_ids[index]
+    seg_indices = [i for i, flag in enumerate(annot["keep_status"]) if flag]
+    keys = ["non_speech_segments", "silence_prob", "audiomae_on_audioset", "duration"]
+    for k in keys:
+        annot[k] = [x for i, x in enumerate(annot[k]) if i in seg_indices]
+    del annot["keep_status"]
+    labels = [
+        summarize_classification_probs(
+            annot["silence_prob"][i], annot["audiomae_on_audioset"][i]
+        ) for i in range(len(annot["non_speech_segments"]))
+    ]
+    segments, durations = annot["non_speech_segments"], annot["duration"]
+    movie = annot["title"]
+
 
-    row = st.session_state.subdf.iloc[index].to_dict()
-    if use_local:
-        clip_paths, labels, segments, durations = process_sample(row)
-    else:
-        annot = load_json(row["annot_filtered"])
-        seg_indices = [i for i, flag in enumerate(annot["keep_status"]) if flag]
-        keys = ["non_speech_segments", "silence_prob", "audiomae_on_audioset", "duration"]
-        for k in keys:
-            annot[k] = [x for i, x in enumerate(annot[k]) if i in seg_indices]
-        del annot["keep_status"]
-        labels = [
-            summarize_classification_probs(
-                annot["silence_prob"][i], annot["audiomae_on_audioset"][i]
-            ) for i in range(len(annot["non_speech_segments"]))
-        ]
-        segments, durations = annot["non_speech_segments"], annot["duration"]
-        clip_paths = [f"https://www.youtube.com/watch?v={row['videoid']}"] * len(segments)
-
-    # Make a grid of videos and captions in streamlit
-    videos = clip_paths
-    video_id = row["videoid"]
-    movie = row["title"]
     st.markdown(f"Showing Foley segments from a clip in movie: **{movie}**")
 
     # Create a grid of videos
     grid = make_grid(3, 3)
 
     # Add videos to the grid
-    for idx in range(0, min(len(videos), 9)):
+    for idx in range(0, min(len(segments), 9)):
         i, j = idx // 3, idx % 3
     
         start, end = segments[idx]
         duration = durations[idx]
 
         grid[i][j].caption(f"Segment duration: {duration}")
-        if not use_local:
-            url = f"https://www.youtube.com/embed/{video_id}?start={int(start)}&end={int(end)}"
-            html_code = f"""
-            <iframe height="320" width="420" src="{url}" frameborder="0" allowfullscreen></iframe>
-            """
-            grid[i][j].markdown(html_code, unsafe_allow_html=True)
-        else:
-            grid[i][j].video(videos[idx])
+        url = f"https://www.youtube.com/embed/{video_id}?start={int(start)}&end={int(end)}"
+        html_code = f"""
+        <iframe height="320" width="420" src="{url}" frameborder="0" allowfullscreen></iframe>
+        """
+        grid[i][j].markdown(html_code, unsafe_allow_html=True)
         grid[i][j].caption(f"{labels[idx]}")
-    
+
     
     st.markdown("##### Some stats")
-    st.write(f"Total number of unique clips: {len(st.session_state.subdf)}")
-    st.write(f"Total number of foley segments: {st.session_state.num_foley}")
+    st.write(f"Total number of unique clips: {len(st.session_state.data)}")
+    st.write("Total number of foley segments: {}".format(st.session_state.num_foley_segments))