Update app.py

app.py

@@ -3,6 +3,8 @@ import os
 import random
 import numpy as np
 import gdown
+from datetime import datetime  
+from csv import writer
 
 from datasets import load_dataset
 from hfserver import HuggingFaceDatasetSaver, HuggingFaceDatasetJSONSaver
@@ -32,7 +34,7 @@ def nan():
 
 FORMAT = ['mp4', 'gif'][1]
 
-def update(data_folder='videos'):
+def update(user_choice, data_folder='videos'):
     # data_folder='videos'
     envs = parse_envs()   
     env_name = envs[random.randint(0, len(envs)-1)]
@@ -46,7 +48,18 @@ def update(data_folder='videos'):
     selected_video_ids = np.random.choice(len(video_files), 2, replace=False)
     left = video_files[selected_video_ids[0]]
     right = video_files[selected_video_ids[1]]
-    print(env_name, left, right)
+
+    # log
+    current_time = datetime.datetime.now()
+    time_stamp = current_time.timestamp()
+    date_time = datetime.fromtimestamp(time_stamp)
+    str_date_time = date_time.strftime("%d-%m-%Y-%H-%M-%S")
+    print(env_name, user_choice, left, right, str_date_time)
+    info = [env_name, user_choice, left, right]
+    with open('data.csv', 'a') as file:
+        writer_object = writer(file)
+        writer_object.writerow(info)
+        file.close()
     return left, right
 
 # def update(left, right):
@@ -112,7 +125,7 @@ def build_interface(iter=3, data_folder='./videos'):
         callback.setup([user_choice, left, right], "flagged_data_points")
         
         btn1.click(fn=replay, inputs=[left, right], outputs=[left, right])
-        btn2.click(fn=update, inputs=None, outputs=[left, right])
+        btn2.click(fn=update, inputs=[user_choice], outputs=[left, right])
 
         # We can choose which components to flag -- in this case, we'll flag all of them
         btn2.click(lambda *args: callback.flag(args), [user_choice, left, right], None, preprocess=False)