download files from drive

app.py

@@ -1,9 +1,11 @@
 import os
 import subprocess
+import glob
 
 import streamlit as st
 
 from utils import get_configs, get_display_names, get_path_for_viz, get_video_height, get_text_str
+from gdrive import download_file
 
 # st.header("EVREAL  - Event-based Video Reconstruction Evaluation and Analysis Library")
 #
@@ -19,8 +21,6 @@ from utils import get_configs, get_display_names, get_path_for_viz, get_video_he
 
 st.title("Result Analysis Tool")
 
-data_base_path = "/home/bercan/ebv/evreal_data"
-
 font_path = "font/Ubuntu-B.ttf"
 
 dataset_cfg_path = os.path.join("cfg", "dataset")
@@ -105,6 +105,11 @@ crop_str = "crop=trunc(iw/2)*2:trunc(ih/2)*2"
 pad_str = "pad=ceil(iw/2)*2+{}:ceil(ih/2)*2+{}:{}:{}".format(padding*2, padding*2, padding, padding)
 num_elements = num_rows * num_cols
 
+# remove previous temp data
+files = glob.glob('temp_data/temp_*.mp4')
+for f in files:
+    os.remove(f)
+
 w = selected_dataset["width"]
 h = selected_dataset["height"]
 input_filter_parts = []
@@ -123,17 +128,21 @@ for row_idx in range(num_rows):
         cur_model = selected_models[col_idx]
         if cur_model['name'] == "gt":
             if row_idx < len(gt_viz):
-                video_path = get_path_for_viz(data_base_path, selected_dataset, selected_sequence, cur_model, gt_viz[row_idx])
-                if not os.path.isfile(video_path):
-                    raise ValueError("Video path does not exist: " + video_path)
+                video_path = get_path_for_viz(selected_dataset, selected_sequence, cur_model, gt_viz[row_idx])
+                local_video_path = os.path.join("temp_data", "temp_{}.mp4".format(vid_idx))
+                download_file(video_path, local_video_path)
+                if not os.path.isfile(local_video_path):
+                    raise ValueError("Failed to download video: " + video_path)
                 gt_viz_indices.append(vid_idx)
             else:
                 continue
         else:
             if row_idx < len(model_viz):
-                video_path = get_path_for_viz(data_base_path, selected_dataset, selected_sequence, cur_model, model_viz[row_idx])
-                if not os.path.isfile(video_path):
-                    raise ValueError("Video path does not exist: " + video_path)
+                video_path = get_path_for_viz(selected_dataset, selected_sequence, cur_model, model_viz[row_idx])
+                local_video_path = os.path.join("temp_data", "temp_{}.mp4".format(vid_idx))
+                download_file(video_path, local_video_path)
+                if not os.path.isfile(local_video_path):
+                    raise ValueError("Failed to download video: " + video_path)
             else:
                 continue
 
@@ -155,7 +164,7 @@ for row_idx in range(num_rows):
         input_filter_parts.append(input_filter_part)
         xstack_input_part = "[v{}]".format(vid_idx)
         xstack_input_parts.append(xstack_input_part)
-        video_paths.append(video_path)
+        video_paths.append(local_video_path)
         if row_idx == 0 or col_idx > 0:
             layout_w_parts = [str(left_pad)] + ["w{}".format(i) for i in range(col_idx)]
             layout_w = "+".join(layout_w_parts)

gdrive.py

@@ -0,0 +1,61 @@
+import os
+
+from pydrive2.auth import GoogleAuth
+from pydrive2.drive import GoogleDrive
+from oauth2client.service_account import ServiceAccountCredentials
+
+
+def login_with_service_account():
+    scope = ["https://www.googleapis.com/auth/drive"]
+    gauth = GoogleAuth()
+    gauth.auth_method = 'service'
+    gauth.credentials = ServiceAccountCredentials.from_json_keyfile_name('evreal-bf06ca1ac056.json', scope)
+    return gauth
+
+
+gauth = login_with_service_account()
+drive = GoogleDrive(gauth)
+
+
+def get_file_id_inside_folder(drive, folder_id, file_name):
+    file_list = drive.ListFile({'q': "'{}' in parents and trashed=false".format(folder_id)}).GetList()
+    for file in file_list:
+        if file['title'] == file_name:
+            return file['id']
+    return None
+
+
+def get_file_id_for_file_path(drive, file_path):
+    # file_parts = os.path.split(file_path)
+    path = os.path.normpath(file_path)
+    file_parts = path.split(os.sep)
+    file_id = '1-5zmRu3Od9s2dYu1S01u4aC2ihAmo1K2'
+    for file_part in file_parts:
+        if file_part == '':
+            continue
+        file_id = get_file_id_inside_folder(drive, file_id, file_part)
+        if file_id is None:
+            return None
+
+    return file_id
+
+
+def download_file(file_path, local_path):
+    file_id = get_file_id_for_file_path(drive, file_path)
+    if file_id is None:
+        return False
+    file = drive.CreateFile({'id': file_id})
+    os.makedirs(os.path.dirname(local_path), exist_ok=True)
+    file.GetContentFile(local_path)
+    return True
+
+
+# test if main
+if __name__ == "__main__":
+    drive_file_path = "0000_000000/ECD_calibration/videos/recon.mp4"
+    # drive_file_path = "test.txt"
+    local_file_path = os.path.join("temp_data", drive_file_path)
+    # gauth = login_with_service_account()
+    # drive = GoogleDrive(gauth)
+    os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+    download_file(drive_file_path, local_file_path)

requirements.txt

@@ -1 +1,2 @@
 ffmpeg-python
+PyDrive2

utils.py

@@ -31,9 +31,9 @@ def get_display_names(configs):
     return display_names
 
 
-def get_path_for_viz(base_path, dataset, sequence, model, viz):
+def get_path_for_viz(dataset, sequence, model, viz):
     dat_seq = "_".join([dataset["name"], sequence])
-    video_path = os.path.join(base_path, model['model_id'], dat_seq, "videos", viz["name"] + ".mp4")
+    video_path = os.path.join(model['model_id'], dat_seq, "videos", viz["name"] + ".mp4")
     return video_path