Spaces:
Running
Running
| import gradio as gr | |
| import numpy as np | |
| import zipfile | |
| import imageio | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| from utils import read_video, frame_sampling | |
| from utils import num_frames, patch_size, input_size | |
| from labels import K400_label_map, SSv2_label_map | |
| LABEL_MAPS = { | |
| 'K400': K400_label_map, | |
| 'SSv2': SSv2_label_map, | |
| } | |
| ALL_MODELS = [ | |
| 'TFUniFormerV2_K400_K710_L14_16x224', | |
| 'TFUniFormerV2_SSV2_B16_16x224', | |
| ] | |
| sample_example = [ | |
| ["examples/k400.mp4", ALL_MODELS[0]], | |
| ["examples/ssv2.mp4", ALL_MODELS[1]], | |
| ] | |
| def get_model(model_type): | |
| model_path = keras.utils.get_file( | |
| origin=f'https://github.com/innat/UniFormerV2/releases/download/v1.1/{model_type}.zip', | |
| ) | |
| with zipfile.ZipFile(model_path, 'r') as zip_ref: | |
| zip_ref.extractall('./') | |
| model = keras.models.load_model(model_type) | |
| if 'K400' in model_type: | |
| data_type = 'K400' | |
| else: | |
| data_type = 'SSv2' | |
| label_map = LABEL_MAPS.get(data_type) | |
| label_map = {v: k for k, v in label_map.items()} | |
| return model, label_map | |
| def inference(video_file, model_type): | |
| # get sample data | |
| container = read_video(video_file) | |
| frames = frame_sampling(container, num_frames=num_frames) | |
| # get models | |
| model, label_map = get_model(model_type) | |
| model.trainable = False | |
| # inference on model | |
| outputs = model(frames[None, ...], training=False) | |
| probabilities = tf.nn.softmax(outputs).numpy().squeeze(0) | |
| confidences = { | |
| label_map[i]: float(probabilities[i]) for i in np.argsort(probabilities)[::-1] | |
| } | |
| return confidences | |
| def main(): | |
| iface = gr.Interface( | |
| fn=inference, | |
| inputs=[ | |
| gr.Video(type="file", label="Input Video"), | |
| gr.Dropdown( | |
| choices=ALL_MODELS, | |
| label="Model" | |
| ) | |
| ], | |
| outputs=gr.Label(num_top_classes=3, label='scores'), | |
| examples=sample_example, | |
| title="UniFormerV2: Spatiotemporal Learning.", | |
| description="Keras reimplementation of <a href='https://github.com/innat/UniFormerV2'>UniFormerV2</a> is presented here." | |
| ) | |
| iface.launch() | |
| if __name__ == '__main__': | |
| main() |