quantumiracle-git/OpenBiDexHand
5
1import gradio as gr2import os3import random4import numpy as np5import pandas as pd6import gdown7import base648from time import gmtime, strftime9from csv import writer10import json11import zipfile12from os import listdir13from os.path import isfile, join, isdir14from datasets import load_dataset15from hfserver import HuggingFaceDatasetSaver, HuggingFaceDatasetJSONSaver16 17ENVS = ['ShadowHand', 'ShadowHandCatchAbreast', 'ShadowHandOver', 'ShadowHandBlockStack', 'ShadowHandCatchUnderarm',18'ShadowHandCatchOver2Underarm', 'ShadowHandBottleCap', 'ShadowHandLiftUnderarm', 'ShadowHandTwoCatchUnderarm',19'ShadowHandDoorOpenInward', 'ShadowHandDoorOpenOutward', 'ShadowHandDoorCloseInward', 'ShadowHandDoorCloseOutward',20'ShadowHandPushBlock', 'ShadowHandKettle', 21'ShadowHandScissors', 'ShadowHandPen', 'ShadowHandSwingCup', 'ShadowHandGraspAndPlace', 'ShadowHandSwitch']22 23# download data from huggingface dataset24# dataset = load_dataset("quantumiracle-git/robotinder-data")25# os.remove('.git/hooks/pre-push') # https://github.com/git-lfs/git-lfs/issues/85326LOAD_DATA_GOOGLE_DRIVE = False27 28if LOAD_DATA_GOOGLE_DRIVE: # download data from google drive29 # url = 'https://drive.google.com/drive/folders/1JuNQS4R7axTezWj1x4KRAuRt_L26ApxA?usp=sharing' # './processed/' folder in google drive30 # url = 'https://drive.google.com/drive/folders/1o8Q9eX-J7F326zv4g2MZWlzR46uVkUF2?usp=sharing' # './processed_zip/' folder in google drive31 # url = 'https://drive.google.com/drive/folders/1ZWgpPiZwnWfwlwta8Tu-Jtu2HsS7HAEa?usp=share_link' # './filter_processed_zip/' folder in google drive32 # url = 'https://drive.google.com/drive/folders/1ROkuX6rQpyK7vLqF5fL2mggKiMCdKSuY?usp=share_link' # './split_processed_zip/' folder in google drive33 34 # output = './'35 # id = url.split('/')[-1]36 # os.system(f"gdown --id {id} -O {output} --folder --no-cookies --remaining-ok")37 # # VIDEO_PATH = 'processed_zip'38 # # VIDEO_PATH = 'filter_processed_zip'39 # VIDEO_PATH = 'split_processed_zip'40 41 # # unzip the zip files to the same location and delete zip files42 # path_to_zip_file = VIDEO_PATH43 # zip_files = [join(path_to_zip_file, f) for f in listdir(path_to_zip_file)]44 # for f in zip_files:45 # if f.endswith(".zip"):46 # directory_to_extract_to = path_to_zip_file # extracted file itself contains a folder47 # print(f'extract data {f} to {directory_to_extract_to}')48 # with zipfile.ZipFile(f, 'r') as zip_ref:49 # zip_ref.extractall(directory_to_extract_to)50 # os.remove(f)51 52 ### multiple urls to handle the retrieve error53 # urls = [54 # 'https://drive.google.com/drive/folders/1BbQe4XtcsalsvwGVLW9jWCkr-ln5pvyf?usp=share_link', # './filter_processed_zip/1' folder in google drive55 # 'https://drive.google.com/drive/folders/1saUTUuObPhMJFguc2J_O0K5woCJjYHci?usp=share_link', # './filter_processed_zip/2' folder in google drive56 # 'https://drive.google.com/drive/folders/1Kh9_E28-RH8g8EP1V3DhGI7KRs9LB7YJ?usp=share_link', # './filter_processed_zip/3' folder in google drive57 # 'https://drive.google.com/drive/folders/1oE75Dz6hxtaSpNhjD22PmQfgQ-PjnEc0?usp=share_link', # './filter_processed_zip/4' folder in google drive58 # 'https://drive.google.com/drive/folders/1XSPEKFqNHpXdLho-bnkT6FZZXssW8JkC?usp=share_link', # './filter_processed_zip/5' folder in google drive59 # 'https://drive.google.com/drive/folders/1XwjAHqR7kF1uSyZZIydQMoETfdvi0aPD?usp=share_link',60 # 'https://drive.google.com/drive/folders/1TceozOWhLsbqP-w-RkforjAVo1M2zsRP?usp=share_link',61 # 'https://drive.google.com/drive/folders/1zAP9eDSW5Eh_isACuZJadXcFaJNqEM9u?usp=share_link',62 # 'https://drive.google.com/drive/folders/1oK8fyF9A3Pv5JubvrQMjTE9n66vYlyZN?usp=share_link',63 # 'https://drive.google.com/drive/folders/1cezGNjlM0ONMM6C0N_PbZVCGsTyVSR0w?usp=share_link',64 # ]65 66 urls = [67 'https://drive.google.com/drive/folders/1SF5jQ7HakO3lFXBon57VP83-AwfnrM3F?usp=share_link', # './split_processed_zip/1' folder in google drive68 'https://drive.google.com/drive/folders/13WuS6ow6sm7ws7A5xzCEhR-2XX_YiIu5?usp=share_link', # './split_processed_zip/2' folder in google drive69 'https://drive.google.com/drive/folders/1GWLffJDOyLkubF2C03UFcB7iFpzy1aDy?usp=share_link', # './split_processed_zip/3' folder in google drive70 'https://drive.google.com/drive/folders/1UKAntA7WliD84AUhRN224PkW4vt9agZW?usp=share_link', # './split_processed_zip/4' folder in google drive71 'https://drive.google.com/drive/folders/11cCQw3qb1vJbviVPfBnOVWVzD_VzHdWs?usp=share_link', # './split_processed_zip/5' folder in google drive72 'https://drive.google.com/drive/folders/1Wvy604wCxEdXAwE7r3sE0L0ieXvM__u8?usp=share_link',73 'https://drive.google.com/drive/folders/1BTv_pMTNGm7m3hD65IgBrX880v-rLIaf?usp=share_link',74 'https://drive.google.com/drive/folders/12x7F11ln2VQkqi8-Mu3kng74eLgifM0N?usp=share_link',75 'https://drive.google.com/drive/folders/1OWkOul2CCrqynqpt44Fu1CBxzNNfOFE2?usp=share_link',76 'https://drive.google.com/drive/folders/1ukwsfrbSEqCBNmRSuAYvYBHijWCQh2OU?usp=share_link',77 'https://drive.google.com/drive/folders/1EO7zumR6sVfsWQWCS6zfNs5WuO2Se6WX?usp=share_link',78 'https://drive.google.com/drive/folders/1aw0iBWvvZiSKng0ejRK8xbNoHLVUFCFu?usp=share_link',79 'https://drive.google.com/drive/folders/1szIcxlVyT5WJtzpqYWYlue0n82A6-xtk?usp=share_link',80 ]81 82 output = './'83 # VIDEO_PATH = 'processed_zip'84 # VIDEO_PATH = 'filter_processed_zip'85 VIDEO_PATH = 'split_processed_zip'86 for i, url in enumerate(urls):87 id = url.split('/')[-1]88 os.system(f"gdown --id {id} -O {output} --folder --no-cookies --remaining-ok")89 90 # unzip the zip files to the same location and delete zip files91 path_to_zip_file = str(i+1)92 zip_files = [join(path_to_zip_file, f) for f in listdir(path_to_zip_file)]93 for f in zip_files:94 if f.endswith(".zip"):95 directory_to_extract_to = VIDEO_PATH # extracted file itself contains a folder96 print(f'extract data {f} to {directory_to_extract_to}')97 with zipfile.ZipFile(f, 'r') as zip_ref:98 zip_ref.extractall(directory_to_extract_to)99 os.remove(f)100 101else:102 VIDEO_PATH = 'processed-data'103 path_to_zip_file = VIDEO_PATH104 zip_files = [join(path_to_zip_file, f) for f in listdir(path_to_zip_file)]105 for f in zip_files:106 if f.endswith(".zip"):107 directory_to_extract_to = path_to_zip_file # extracted file itself contains a folder108 print(f'extract data {f} to {directory_to_extract_to}')109 with zipfile.ZipFile(f, 'r') as zip_ref:110 zip_ref.extractall(directory_to_extract_to)111 os.remove(f)112 113# for test only114# else: # local data115# VIDEO_PATH = 'robotinder-data'116 117VIDEO_INFO = os.path.join(VIDEO_PATH, 'video_info.json')118 119def inference(video_path):120 # for displaying mp4 with autoplay on Gradio121 with open(video_path, "rb") as f:122 data = f.read()123 b64 = base64.b64encode(data).decode()124 html = (125 f"""126 <video controls autoplay muted loop>127 <source src="data:video/mp4;base64,{b64}" type="video/mp4">128 </video> 129 """130 )131 return html132 133def video_identity(video):134 return video135 136def nan():137 return None138 139FORMAT = ['mp4', 'gif'][0]140 141def get_huggingface_dataset():142 try:143 import huggingface_hub144 except (ImportError, ModuleNotFoundError):145 raise ImportError(146 "Package `huggingface_hub` not found is needed "147 "for HuggingFaceDatasetSaver. Try 'pip install huggingface_hub'."148 )149 HF_TOKEN = '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' # my HF token150 DATASET_NAME = 'crowdsourced-robotinder-demo'151 FLAGGING_DIR = 'flag/'152 path_to_dataset_repo = huggingface_hub.create_repo(153 repo_id=DATASET_NAME,154 token=HF_TOKEN,155 private=False,156 repo_type="dataset",157 exist_ok=True,158 ) 159 dataset_dir = os.path.join(DATASET_NAME, FLAGGING_DIR)160 repo = huggingface_hub.Repository(161 local_dir=dataset_dir,162 clone_from=path_to_dataset_repo,163 use_auth_token=HF_TOKEN,164 )165 repo.git_pull(lfs=True)166 log_file = os.path.join(dataset_dir, "flag_data.csv")167 return repo, log_file168 169def update(user_choice, user_name, left, right, choose_env, data_folder=VIDEO_PATH, flag_to_huggingface=False):170 global last_left_video_path 171 global last_right_video_path 172 global last_infer_left_video_path173 global last_infer_right_video_path174 175 if flag_to_huggingface: # log176 env_name = str(last_left_video_path).split('/')[1] # 'robotinder-data/ENV_NAME/'177 current_time = strftime("%Y-%m-%d-%H-%M-%S", gmtime())178 info = [env_name, user_choice, last_left_video_path, last_right_video_path, current_time, user_name]179 print(info)180 repo, log_file = get_huggingface_dataset()181 with open(log_file, 'a') as file: # incremental change of the file182 writer_object = writer(file)183 writer_object.writerow(info)184 file.close()185 if int(current_time.split('-')[-2]) % 5 == 0: # push only on certain minutes186 try:187 repo.push_to_hub(commit_message=f"Flagged sample at {current_time}")188 except:189 repo.git_pull(lfs=True) # sync with remote first190 repo.push_to_hub(commit_message=f"Flagged sample at {current_time}")191 if choose_env == 'Random' or choose_env == '': # random or no selection192 envs = get_env_names() 193 env_name = envs[random.randint(0, len(envs)-1)]194 else:195 env_name = choose_env196 # choose video197 left, right = randomly_select_videos(env_name)198 199 last_left_video_path = left200 last_right_video_path = right201 last_infer_left_video_path = inference(left)202 last_infer_right_video_path = inference(right)203 204 return last_infer_left_video_path, last_infer_right_video_path, env_name205 206def replay(left, right): 207 return left, right208 209def parse_envs(folder=VIDEO_PATH, filter=True, MAX_ITER=20000, DEFAULT_ITER=20000):210 """211 return a dict of env_name: video_paths212 """213 files = {}214 if filter:215 df = pd.read_csv('Bidexhands_Video.csv')216 # print(df)217 for env_name in os.listdir(folder):218 env_path = os.path.join(folder, env_name)219 if os.path.isdir(env_path):220 videos = os.listdir(env_path)221 video_files = []222 for video in videos: # video name rule: EnvName_Alg_Seed_Timestamp_Checkpoint_video-episode-EpisodeID223 if video.endswith(f'.{FORMAT}'):224 if filter:225 if len(video.split('_')) < 6:226 print(f'{video} is wrongly named.')227 seed = video.split('_')[2]228 checkpoint = video.split('_')[4]229 try:230 succeed_iteration = df.loc[(df['seed'] == int(seed)) & (df['env_name'] == str(env_name))]['succeed_iteration'].iloc[0]231 except:232 print(f'Env {env_name} with seed {seed} not found in Bidexhands_Video.csv')233 234 if 'unsolved' in succeed_iteration:235 continue236 elif pd.isnull(succeed_iteration):237 min_iter = DEFAULT_ITER238 max_iter = MAX_ITER239 elif '-' in succeed_iteration:240 [min_iter, max_iter] = succeed_iteration.split('-')241 else:242 min_iter = succeed_iteration243 max_iter = MAX_ITER244 245 # check if the checkpoint is in the valid range246 valid_checkpoints = np.arange(int(min_iter), int(max_iter)+1000, 1000)247 if int(checkpoint) not in valid_checkpoints:248 continue249 250 video_path = os.path.join(folder, env_name, video)251 video_files.append(video_path)252 # print(video_path)253 254 files[env_name] = video_files255 256 with open(VIDEO_INFO, 'w') as fp:257 json.dump(files, fp)258 259 return files260 261def get_env_names():262 with open(VIDEO_INFO, 'r') as fp:263 files = json.load(fp)264 return list(files.keys())265 266def randomly_select_videos(env_name):267 # load the parsed video info268 with open(VIDEO_INFO, 'r') as fp:269 files = json.load(fp)270 env_files = files[env_name]271 # randomly choose two videos272 selected_video_ids = np.random.choice(len(env_files), 2, replace=False)273 left_video_path = env_files[selected_video_ids[0]]274 right_video_path = env_files[selected_video_ids[1]]275 return left_video_path, right_video_path276 277def build_interface(iter=3, data_folder=VIDEO_PATH):278 import sys279 import csv280 csv.field_size_limit(sys.maxsize)281 282 HF_TOKEN = os.getenv('HF_TOKEN')283 print(HF_TOKEN)284 HF_TOKEN = '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' # my HF token285 ## hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "crowdsourced-robotinder-demo") # HuggingFace logger instead of local one: https://github.com/gradio-app/gradio/blob/master/gradio/flagging.py286 ## callback = gr.CSVLogger()287 # hf_writer = HuggingFaceDatasetSaver(HF_TOKEN, "crowdsourced-robotinder-demo")288 # callback = hf_writer289 290 # parse the video folder 291 files = parse_envs() 292 293 # build gradio interface294 with gr.Blocks() as demo:295 # gr.Markdown("## Here is <span style=color:cyan>RoboTinder</span>!")296 gr.Markdown("### Select the best robot behaviour in your choice!")297 # some initial values298 env_name = list(files.keys())[random.randint(0, len(files)-1)] # random pick an env 299 with gr.Row():300 str_env_name = gr.Markdown(f"{env_name}")301 302 # choose video303 left_video_path, right_video_path = randomly_select_videos(env_name)304 305 with gr.Row():306 if FORMAT == 'mp4':307 # left = gr.PlayableVideo(left_video_path, label="left_video")308 # right = gr.PlayableVideo(right_video_path, label="right_video")309 310 infer_left_video_path = inference(left_video_path)311 infer_right_video_path = inference(right_video_path)312 left = gr.HTML(infer_left_video_path, label="left_video")313 right = gr.HTML(infer_right_video_path, label="right_video")314 else:315 left = gr.Image(left_video_path, shape=(1024, 768), label="left_video")316 # right = gr.Image(right_video_path).style(height=768, width=1024)317 right = gr.Image(right_video_path, label="right_video")318 319 global last_left_video_path 320 last_left_video_path = left_video_path321 global last_right_video_path 322 last_right_video_path = right_video_path323 324 global last_infer_left_video_path325 last_infer_left_video_path = infer_left_video_path326 global last_infer_right_video_path327 last_infer_right_video_path = infer_right_video_path328 329 # btn1 = gr.Button("Replay")330 user_name = gr.Textbox(label='Your name/email:')331 # user_choice = gr.Radio(["Left", "Right", "Not Sure", "Both Good", "Both Bad"], label="Which one is your favorite?")332 user_choice = gr.Radio(["Left", "Right", "Not Sure"], label="Which one is your favorite?")333 choose_env = gr.Radio(["Random"]+ENVS, label="Choose the next task:")334 btn2 = gr.Button("Next")335 336 # This needs to be called at some point prior to the first call to callback.flag()337 # callback.setup([user_choice, left, right], "flagged_data_points")338 339 # btn1.click(fn=replay, inputs=[left, right], outputs=[left, right])340 btn2.click(fn=update, inputs=[user_choice, user_name, left, right, choose_env], outputs=[left, right, str_env_name])341 342 # We can choose which components to flag -- in this case, we'll flag all of them343 # btn2.click(lambda *args: callback.flag(args), [user_choice, left, right], None, preprocess=False) # not using the gradio flagging anymore344 345 return demo346 347if __name__ == "__main__":348 last_left_video_path = None349 last_right_video_path = None350 351 demo = build_interface()352 # demo.launch(share=True)353 demo.launch(share=False)354