Spaces:
Runtime error
Runtime error
import gradio as gr | |
import numpy as np | |
import zipfile | |
import imageio | |
import tensorflow as tf | |
from tensorflow import keras | |
from utils import TubeMaskingGenerator | |
from utils import read_video, frame_sampling, denormalize, reconstrunction | |
from utils import IMAGENET_MEAN, IMAGENET_STD, num_frames, patch_size, input_size | |
from labels import K400_label_map, SSv2_label_map, UCF_label_map | |
LABEL_MAPS = { | |
'K400': K400_label_map, | |
'SSv2': SSv2_label_map, | |
'UCF' : UCF_label_map | |
} | |
ALL_MODELS = [ | |
'TFVideoMAE_L_K400_16x224', | |
'TFVideoMAE_B_SSv2_16x224', | |
'TFVideoMAE_B_UCF_16x224', | |
] | |
sample_example = [ | |
["examples/k400.mp4", ALL_MODELS[0], 0.9], | |
["examples/ssv2.mp4", ALL_MODELS[1], 0.8], | |
["examples/ucf.mp4", ALL_MODELS[2], 0.7], | |
] | |
def tube_mask_generator(mask_ratio): | |
window_size = ( | |
num_frames // 2, | |
input_size // patch_size[0], | |
input_size // patch_size[1] | |
) | |
tube_mask = TubeMaskingGenerator( | |
input_size=window_size, | |
mask_ratio=mask_ratio | |
) | |
make_bool = tube_mask() | |
bool_masked_pos_tf = tf.constant(make_bool, dtype=tf.int32) | |
bool_masked_pos_tf = tf.expand_dims(bool_masked_pos_tf, axis=0) | |
bool_masked_pos_tf = tf.cast(bool_masked_pos_tf, tf.bool) | |
return bool_masked_pos_tf | |
def get_model(model_type): | |
ft_path = keras.utils.get_file( | |
origin=f'https://github.com/innat/VideoMAE/releases/download/v1.1/{model_type}_FT.zip', | |
) | |
pt_path = keras.utils.get_file( | |
origin=f'https://github.com/innat/VideoMAE/releases/download/v1.1/{model_type}_PT.zip', | |
) | |
with zipfile.ZipFile(ft_path, 'r') as zip_ref: | |
zip_ref.extractall('./') | |
with zipfile.ZipFile(pt_path, 'r') as zip_ref: | |
zip_ref.extractall('./') | |
ft_model = keras.models.load_model(model_type + '_FT') | |
pt_model = keras.models.load_model(model_type + '_PT') | |
if 'K400' in model_type: | |
data_type = 'K400' | |
elif 'SSv2' in model_type: | |
data_type = 'SSv2' | |
else: | |
data_type = 'UCF' | |
label_map = LABEL_MAPS.get(data_type) | |
label_map = {v: k for k, v in label_map.items()} | |
return ft_model, pt_model, label_map | |
def inference(video_file, model_type, mask_ratio): | |
# get sample data | |
container = read_video(video_file) | |
frames = frame_sampling(container, num_frames=num_frames) | |
# get models | |
bool_masked_pos_tf = tube_mask_generator(mask_ratio) | |
ft_model, pt_model, label_map = get_model(model_type) | |
ft_model.trainable = False | |
pt_model.trainable = False | |
# inference on fine-tune model | |
outputs_ft = ft_model(frames[None, ...], training=False) | |
probabilities = tf.nn.softmax(outputs_ft).numpy().squeeze(0) | |
confidences = { | |
label_map[i]: float(probabilities[i]) for i in np.argsort(probabilities)[::-1] | |
} | |
# inference on pre-trained model | |
outputs_pt = pt_model(frames[None, ...], bool_masked_pos_tf, training=False) | |
reconstruct_output, mask = reconstrunction( | |
frames[None, ...], bool_masked_pos_tf, outputs_pt | |
) | |
# post process | |
input_frame = denormalize(frames) | |
input_mask = denormalize(mask[0] * frames) | |
output_frame = denormalize(reconstruct_output) | |
frames = [] | |
for frame_a, frame_b, frame_c in zip(input_frame, input_mask, output_frame): | |
combined_frame = np.hstack([frame_a, frame_b, frame_c]) | |
frames.append(combined_frame) | |
combined_gif = 'combined.gif' | |
imageio.mimsave(combined_gif, frames, duration=300, loop=0) | |
return confidences, combined_gif | |
def main(): | |
iface = gr.Interface( | |
fn=inference, | |
inputs=[ | |
gr.Video(type="file", label="Input Video"), | |
gr.Dropdown( | |
choices=ALL_MODELS, | |
default="TFVideoMAE_L_K400_16x224", | |
label="Model" | |
), | |
gr.Slider( | |
0.5, | |
1.0, | |
step=0.1, | |
default=0.5, | |
label='Mask Ratio' | |
) | |
], | |
outputs=[ | |
gr.Label(num_top_classes=3, label='scores'), | |
gr.Image(type="filepath", label='reconstructed') | |
], | |
examples=sample_example, | |
title="VideoMAE", | |
description="Keras reimplementation of <a href='https://github.com/innat/VideoMAE'>VideoMAE</a> is presented here." | |
) | |
iface.launch() | |
if __name__ == '__main__': | |
main() |