Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| import random | |
| import numpy as np | |
| import pandas as pd | |
| import gdown | |
| import base64 | |
| from time import gmtime, strftime | |
| from csv import writer | |
| import json | |
| import zipfile | |
| from os import listdir | |
| from os.path import isfile, join, isdir | |
| from datasets import load_dataset | |
| from hfserver import HuggingFaceDatasetSaver, HuggingFaceDatasetJSONSaver | |
| ENVS = ['ShadowHand', 'ShadowHandCatchAbreast', 'ShadowHandOver', 'ShadowHandBlockStack', 'ShadowHandCatchUnderarm', | |
| 'ShadowHandCatchOver2Underarm', 'ShadowHandBottleCap', 'ShadowHandLiftUnderarm', 'ShadowHandTwoCatchUnderarm', | |
| 'ShadowHandDoorOpenInward', 'ShadowHandDoorOpenOutward', 'ShadowHandDoorCloseInward', 'ShadowHandDoorCloseOutward', | |
| 'ShadowHandPushBlock', 'ShadowHandKettle', | |
| 'ShadowHandScissors', 'ShadowHandPen', 'ShadowHandSwingCup', 'ShadowHandGraspAndPlace', 'ShadowHandSwitch'] | |
| # download data from huggingface dataset | |
| # dataset = load_dataset("quantumiracle-git/robotinder-data") | |
| # os.remove('.git/hooks/pre-push') # https://github.com/git-lfs/git-lfs/issues/853 | |
| LOAD_DATA_GOOGLE_DRIVE = False | |
| if LOAD_DATA_GOOGLE_DRIVE: # download data from google drive | |
| # url = 'https://drive.google.com/drive/folders/1JuNQS4R7axTezWj1x4KRAuRt_L26ApxA?usp=sharing' # './processed/' folder in google drive | |
| # url = 'https://drive.google.com/drive/folders/1o8Q9eX-J7F326zv4g2MZWlzR46uVkUF2?usp=sharing' # './processed_zip/' folder in google drive | |
| # url = 'https://drive.google.com/drive/folders/1ZWgpPiZwnWfwlwta8Tu-Jtu2HsS7HAEa?usp=share_link' # './filter_processed_zip/' folder in google drive | |
| # url = 'https://drive.google.com/drive/folders/1ROkuX6rQpyK7vLqF5fL2mggKiMCdKSuY?usp=share_link' # './split_processed_zip/' folder in google drive | |
| # output = './' | |
| # id = url.split('/')[-1] | |
| # os.system(f"gdown --id {id} -O {output} --folder --no-cookies --remaining-ok") | |
| # # VIDEO_PATH = 'processed_zip' | |
| # # VIDEO_PATH = 'filter_processed_zip' | |
| # VIDEO_PATH = 'split_processed_zip' | |
| # # unzip the zip files to the same location and delete zip files | |
| # path_to_zip_file = VIDEO_PATH | |
| # zip_files = [join(path_to_zip_file, f) for f in listdir(path_to_zip_file)] | |
| # for f in zip_files: | |
| # if f.endswith(".zip"): | |
| # directory_to_extract_to = path_to_zip_file # extracted file itself contains a folder | |
| # print(f'extract data {f} to {directory_to_extract_to}') | |
| # with zipfile.ZipFile(f, 'r') as zip_ref: | |
| # zip_ref.extractall(directory_to_extract_to) | |
| # os.remove(f) | |
| ### multiple urls to handle the retrieve error | |
| # urls = [ | |
| # 'https://drive.google.com/drive/folders/1BbQe4XtcsalsvwGVLW9jWCkr-ln5pvyf?usp=share_link', # './filter_processed_zip/1' folder in google drive | |
| # 'https://drive.google.com/drive/folders/1saUTUuObPhMJFguc2J_O0K5woCJjYHci?usp=share_link', # './filter_processed_zip/2' folder in google drive | |
| # 'https://drive.google.com/drive/folders/1Kh9_E28-RH8g8EP1V3DhGI7KRs9LB7YJ?usp=share_link', # './filter_processed_zip/3' folder in google drive | |
| # 'https://drive.google.com/drive/folders/1oE75Dz6hxtaSpNhjD22PmQfgQ-PjnEc0?usp=share_link', # './filter_processed_zip/4' folder in google drive | |
| # 'https://drive.google.com/drive/folders/1XSPEKFqNHpXdLho-bnkT6FZZXssW8JkC?usp=share_link', # './filter_processed_zip/5' folder in google drive | |
| # 'https://drive.google.com/drive/folders/1XwjAHqR7kF1uSyZZIydQMoETfdvi0aPD?usp=share_link', | |
| # 'https://drive.google.com/drive/folders/1TceozOWhLsbqP-w-RkforjAVo1M2zsRP?usp=share_link', | |
| # 'https://drive.google.com/drive/folders/1zAP9eDSW5Eh_isACuZJadXcFaJNqEM9u?usp=share_link', | |
| # 'https://drive.google.com/drive/folders/1oK8fyF9A3Pv5JubvrQMjTE9n66vYlyZN?usp=share_link', | |
| # 'https://drive.google.com/drive/folders/1cezGNjlM0ONMM6C0N_PbZVCGsTyVSR0w?usp=share_link', | |
| # ] | |
| urls = [ | |
| 'https://drive.google.com/drive/folders/1SF5jQ7HakO3lFXBon57VP83-AwfnrM3F?usp=share_link', # './split_processed_zip/1' folder in google drive | |
| 'https://drive.google.com/drive/folders/13WuS6ow6sm7ws7A5xzCEhR-2XX_YiIu5?usp=share_link', # './split_processed_zip/2' folder in google drive | |
| 'https://drive.google.com/drive/folders/1GWLffJDOyLkubF2C03UFcB7iFpzy1aDy?usp=share_link', # './split_processed_zip/3' folder in google drive | |
| 'https://drive.google.com/drive/folders/1UKAntA7WliD84AUhRN224PkW4vt9agZW?usp=share_link', # './split_processed_zip/4' folder in google drive | |
| 'https://drive.google.com/drive/folders/11cCQw3qb1vJbviVPfBnOVWVzD_VzHdWs?usp=share_link', # './split_processed_zip/5' folder in google drive | |
| 'https://drive.google.com/drive/folders/1Wvy604wCxEdXAwE7r3sE0L0ieXvM__u8?usp=share_link', | |
| 'https://drive.google.com/drive/folders/1BTv_pMTNGm7m3hD65IgBrX880v-rLIaf?usp=share_link', | |
| 'https://drive.google.com/drive/folders/12x7F11ln2VQkqi8-Mu3kng74eLgifM0N?usp=share_link', | |
| 'https://drive.google.com/drive/folders/1OWkOul2CCrqynqpt44Fu1CBxzNNfOFE2?usp=share_link', | |
| 'https://drive.google.com/drive/folders/1ukwsfrbSEqCBNmRSuAYvYBHijWCQh2OU?usp=share_link', | |
| 'https://drive.google.com/drive/folders/1EO7zumR6sVfsWQWCS6zfNs5WuO2Se6WX?usp=share_link', | |
| 'https://drive.google.com/drive/folders/1aw0iBWvvZiSKng0ejRK8xbNoHLVUFCFu?usp=share_link', | |
| 'https://drive.google.com/drive/folders/1szIcxlVyT5WJtzpqYWYlue0n82A6-xtk?usp=share_link', | |
| ] | |
| output = './' | |
| # VIDEO_PATH = 'processed_zip' | |
| # VIDEO_PATH = 'filter_processed_zip' | |
| VIDEO_PATH = 'split_processed_zip' | |
| for i, url in enumerate(urls): | |
| id = url.split('/')[-1] | |
| os.system(f"gdown --id {id} -O {output} --folder --no-cookies --remaining-ok") | |
| # unzip the zip files to the same location and delete zip files | |
| path_to_zip_file = str(i+1) | |
| zip_files = [join(path_to_zip_file, f) for f in listdir(path_to_zip_file)] | |
| for f in zip_files: | |
| if f.endswith(".zip"): | |
| directory_to_extract_to = VIDEO_PATH # extracted file itself contains a folder | |
| print(f'extract data {f} to {directory_to_extract_to}') | |
| with zipfile.ZipFile(f, 'r') as zip_ref: | |
| zip_ref.extractall(directory_to_extract_to) | |
| os.remove(f) | |
| else: | |
| VIDEO_PATH = 'processed-data' | |
| path_to_zip_file = VIDEO_PATH | |
| zip_files = [join(path_to_zip_file, f) for f in listdir(path_to_zip_file)] | |
| for f in zip_files: | |
| if f.endswith(".zip"): | |
| directory_to_extract_to = path_to_zip_file # extracted file itself contains a folder | |
| print(f'extract data {f} to {directory_to_extract_to}') | |
| with zipfile.ZipFile(f, 'r') as zip_ref: | |
| zip_ref.extractall(directory_to_extract_to) | |
| os.remove(f) | |
| # for test only | |
| # else: # local data | |
| # VIDEO_PATH = 'robotinder-data' | |
| VIDEO_INFO = os.path.join(VIDEO_PATH, 'video_info.json') | |
| def inference(video_path): | |
| # for displaying mp4 with autoplay on Gradio | |
| with open(video_path, "rb") as f: | |
| data = f.read() | |
| b64 = base64.b64encode(data).decode() | |
| html = ( | |
| f""" | |
| <video controls autoplay muted loop> | |
| <source src="data:video/mp4;base64,{b64}" type="video/mp4"> | |
| </video> | |
| """ | |
| ) | |
| return html | |
| def video_identity(video): | |
| return video | |
| def nan(): | |
| return None | |
| FORMAT = ['mp4', 'gif'][0] | |
| def get_huggingface_dataset(): | |
| try: | |
| import huggingface_hub | |
| except (ImportError, ModuleNotFoundError): | |
| raise ImportError( | |
| "Package `huggingface_hub` not found is needed " | |
| "for HuggingFaceDatasetSaver. Try 'pip install huggingface_hub'." | |
| ) | |
| HF_TOKEN = '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' # my HF token | |
| DATASET_NAME = 'crowdsourced-robotinder-demo' | |
| FLAGGING_DIR = 'flag/' | |
| path_to_dataset_repo = huggingface_hub.create_repo( | |
| repo_id=DATASET_NAME, | |
| token=HF_TOKEN, | |
| private=False, | |
| repo_type="dataset", | |
| exist_ok=True, | |
| ) | |
| dataset_dir = os.path.join(DATASET_NAME, FLAGGING_DIR) | |
| repo = huggingface_hub.Repository( | |
| local_dir=dataset_dir, | |
| clone_from=path_to_dataset_repo, | |
| use_auth_token=HF_TOKEN, | |
| ) | |
| repo.git_pull(lfs=True) | |
| log_file = os.path.join(dataset_dir, "flag_data.csv") | |
| return repo, log_file | |
| def update(user_choice, user_name, left, right, choose_env, data_folder=VIDEO_PATH, flag_to_huggingface=False): | |
| global last_left_video_path | |
| global last_right_video_path | |
| global last_infer_left_video_path | |
| global last_infer_right_video_path | |
| if flag_to_huggingface: # log | |
| env_name = str(last_left_video_path).split('/')[1] # 'robotinder-data/ENV_NAME/' | |
| current_time = strftime("%Y-%m-%d-%H-%M-%S", gmtime()) | |
| info = [env_name, user_choice, last_left_video_path, last_right_video_path, current_time, user_name] | |
| print(info) | |
| repo, log_file = get_huggingface_dataset() | |
| with open(log_file, 'a') as file: # incremental change of the file | |
| writer_object = writer(file) | |
| writer_object.writerow(info) | |
| file.close() | |
| if int(current_time.split('-')[-2]) % 5 == 0: # push only on certain minutes | |
| try: | |
| repo.push_to_hub(commit_message=f"Flagged sample at {current_time}") | |
| except: | |
| repo.git_pull(lfs=True) # sync with remote first | |
| repo.push_to_hub(commit_message=f"Flagged sample at {current_time}") | |
| if choose_env == 'Random' or choose_env == '': # random or no selection | |
| envs = get_env_names() | |
| env_name = envs[random.randint(0, len(envs)-1)] | |
| else: | |
| env_name = choose_env | |
| # choose video | |
| left, right = randomly_select_videos(env_name) | |
| last_left_video_path = left | |
| last_right_video_path = right | |
| last_infer_left_video_path = inference(left) | |
| last_infer_right_video_path = inference(right) | |
| return last_infer_left_video_path, last_infer_right_video_path, env_name | |
| def replay(left, right): | |
| return left, right | |
| def parse_envs(folder=VIDEO_PATH, filter=True, MAX_ITER=20000, DEFAULT_ITER=20000): | |
| """ | |
| return a dict of env_name: video_paths | |
| """ | |
| files = {} | |
| if filter: | |
| df = pd.read_csv('Bidexhands_Video.csv') | |
| # print(df) | |
| for env_name in os.listdir(folder): | |
| env_path = os.path.join(folder, env_name) | |
| if os.path.isdir(env_path): | |
| videos = os.listdir(env_path) | |
| video_files = [] | |
| for video in videos: # video name rule: EnvName_Alg_Seed_Timestamp_Checkpoint_video-episode-EpisodeID | |
| if video.endswith(f'.{FORMAT}'): | |
| if filter: | |
| if len(video.split('_')) < 6: | |
| print(f'{video} is wrongly named.') | |
| seed = video.split('_')[2] | |
| checkpoint = video.split('_')[4] | |
| try: | |
| succeed_iteration = df.loc[(df['seed'] == int(seed)) & (df['env_name'] == str(env_name))]['succeed_iteration'].iloc[0] | |
| except: | |
| print(f'Env {env_name} with seed {seed} not found in Bidexhands_Video.csv') | |
| if 'unsolved' in succeed_iteration: | |
| continue | |
| elif pd.isnull(succeed_iteration): | |
| min_iter = DEFAULT_ITER | |
| max_iter = MAX_ITER | |
| elif '-' in succeed_iteration: | |
| [min_iter, max_iter] = succeed_iteration.split('-') | |
| else: | |
| min_iter = succeed_iteration | |
| max_iter = MAX_ITER | |
| # check if the checkpoint is in the valid range | |
| valid_checkpoints = np.arange(int(min_iter), int(max_iter)+1000, 1000) | |
| if int(checkpoint) not in valid_checkpoints: | |
| continue | |
| video_path = os.path.join(folder, env_name, video) | |
| video_files.append(video_path) | |
| # print(video_path) | |
| files[env_name] = video_files | |
| with open(VIDEO_INFO, 'w') as fp: | |
| json.dump(files, fp) | |
| return files | |
| def get_env_names(): | |
| with open(VIDEO_INFO, 'r') as fp: | |
| files = json.load(fp) | |
| return list(files.keys()) | |
| def randomly_select_videos(env_name): | |
| # load the parsed video info | |
| with open(VIDEO_INFO, 'r') as fp: | |
| files = json.load(fp) | |
| env_files = files[env_name] | |
| # randomly choose two videos | |
| selected_video_ids = np.random.choice(len(env_files), 2, replace=False) | |
| left_video_path = env_files[selected_video_ids[0]] | |
| right_video_path = env_files[selected_video_ids[1]] | |
| return left_video_path, right_video_path | |
| def build_interface(iter=3, data_folder=VIDEO_PATH): | |
| import sys | |
| import csv | |
| csv.field_size_limit(sys.maxsize) | |
| HF_TOKEN = os.getenv('HF_TOKEN') | |
| print(HF_TOKEN) | |
| HF_TOKEN = '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' # my HF token | |
| ## 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.py | |
| ## callback = gr.CSVLogger() | |
| # hf_writer = HuggingFaceDatasetSaver(HF_TOKEN, "crowdsourced-robotinder-demo") | |
| # callback = hf_writer | |
| # parse the video folder | |
| files = parse_envs() | |
| # build gradio interface | |
| with gr.Blocks() as demo: | |
| # gr.Markdown("## Here is <span style=color:cyan>RoboTinder</span>!") | |
| gr.Markdown("### Select the best robot behaviour in your choice!") | |
| # some initial values | |
| env_name = list(files.keys())[random.randint(0, len(files)-1)] # random pick an env | |
| with gr.Row(): | |
| str_env_name = gr.Markdown(f"{env_name}") | |
| # choose video | |
| left_video_path, right_video_path = randomly_select_videos(env_name) | |
| with gr.Row(): | |
| if FORMAT == 'mp4': | |
| # left = gr.PlayableVideo(left_video_path, label="left_video") | |
| # right = gr.PlayableVideo(right_video_path, label="right_video") | |
| infer_left_video_path = inference(left_video_path) | |
| infer_right_video_path = inference(right_video_path) | |
| left = gr.HTML(infer_left_video_path, label="left_video") | |
| right = gr.HTML(infer_right_video_path, label="right_video") | |
| else: | |
| left = gr.Image(left_video_path, shape=(1024, 768), label="left_video") | |
| # right = gr.Image(right_video_path).style(height=768, width=1024) | |
| right = gr.Image(right_video_path, label="right_video") | |
| global last_left_video_path | |
| last_left_video_path = left_video_path | |
| global last_right_video_path | |
| last_right_video_path = right_video_path | |
| global last_infer_left_video_path | |
| last_infer_left_video_path = infer_left_video_path | |
| global last_infer_right_video_path | |
| last_infer_right_video_path = infer_right_video_path | |
| # btn1 = gr.Button("Replay") | |
| user_name = gr.Textbox(label='Your name/email:') | |
| # user_choice = gr.Radio(["Left", "Right", "Not Sure", "Both Good", "Both Bad"], label="Which one is your favorite?") | |
| user_choice = gr.Radio(["Left", "Right", "Not Sure"], label="Which one is your favorite?") | |
| choose_env = gr.Radio(["Random"]+ENVS, label="Choose the next task:") | |
| btn2 = gr.Button("Next") | |
| # This needs to be called at some point prior to the first call to callback.flag() | |
| # callback.setup([user_choice, left, right], "flagged_data_points") | |
| # btn1.click(fn=replay, inputs=[left, right], outputs=[left, right]) | |
| btn2.click(fn=update, inputs=[user_choice, user_name, left, right, choose_env], outputs=[left, right, str_env_name]) | |
| # 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) # not using the gradio flagging anymore | |
| return demo | |
| if __name__ == "__main__": | |
| last_left_video_path = None | |
| last_right_video_path = None | |
| demo = build_interface() | |
| # demo.launch(share=True) | |
| demo.launch(share=False) | |