using hf datasets for videos

app.py

@@ -1,8 +1,8 @@
 import csv
 import os
+import pickle
 import random
 import sys
-import pickle
 from collections import Counter
 from glob import glob
 
@@ -10,23 +10,17 @@ import clip
 import gdown
 import gradio as gr
 import numpy as np
+import psutil
 import torch
 import torchvision
+from datasets import load_dataset
 from tqdm import tqdm
 
 from SimSearch import FaissCosineNeighbors
 
 csv.field_size_limit(sys.maxsize)
 
-
-# DOWNLOAD THE DATASET and Files
-gdown.cached_download(
-    url="https://static.taesiri.com/gamephysics/GTAV-Videos.zip",
-    path="./GTAV-Videos.zip",
-    quiet=False,
-    md5="e961093bc032f579de060ed65564b4c3",
-)
-
+# Download Embeddings
 gdown.cached_download(
     url="https://static.taesiri.com/gamephysics/mini-GTA-V-Embeddings.zip",
     path="./GTA-V-Embeddings.zip",
@@ -34,10 +28,6 @@ gdown.cached_download(
     md5="b1228503d5a89eef7e35e2cbf86b2fc0",
 )
 
-# EXTRACT
-torchvision.datasets.utils.extract_archive(
-    from_path="GTAV-Videos.zip", to_path="Videos/", remove_finished=False
-)
 # EXTRACT
 torchvision.datasets.utils.extract_archive(
     from_path="GTA-V-Embeddings.zip",
@@ -45,9 +35,23 @@ torchvision.datasets.utils.extract_archive(
     remove_finished=False,
 )
 
+# Load videos from Dataset
+gta_v_videos = load_dataset("taesiri/GamePhysics_Grand_Theft_Auto_V")
+post_id_to_video_path = {
+    os.path.splitext(os.path.basename(x))[0]: x
+    for x in gta_v_videos["Grand_Theft_Auto_V"][:]["video_file_path"]
+}
 # Initialize CLIP model
 clip.available_models()
 
+
+# Log runtime environment info
+def log_runtime_information():
+    print(f"CPU Count: {psutil.cpu_count()}")
+    print(f"Virtual Memory: {psutil.virtual_memory()}")
+    print(f"Swap Memory: {psutil.swap_memory()}")
+
+
 # # Searcher
 class GamePhysicsSearcher:
     def __init__(self, CLIP_MODEL, GAME_NAME, EMBEDDING_PATH="./Embeddings/VIT32/"):
@@ -113,8 +117,7 @@ class GamePhysicsSearcher:
             topKs = list(self.f7(nearest_data_points[0]))[:k]
 
         video_filename = [
-            f"./Videos/{self.GAME_NAME}/"
-            + self.class_id_to_file[x].replace("npy", "mp4")
+            post_id_to_video_path[self.class_id_to_file[x].replace(".npy", "")]
             for x in topKs
         ]
 
@@ -239,4 +242,5 @@ def main():
 
 
 if __name__ == "__main__":
+    log_runtime_information()
     main()

requirements.txt

@@ -59,4 +59,5 @@ typing_extensions==4.0.1
 urllib3==1.26.8
 wcwidth==0.2.5
 Werkzeug==2.0.2
+psuti==5.9.1
 gdown