Spaces:
Build error
Build error
File size: 3,947 Bytes
ccdf9bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import copy
import torch
import numpy as np
import gradio as gr
from spoter_mod.skeleton_extractor import obtain_pose_data
from spoter_mod.normalization.body_normalization import normalize_single_dict as normalize_single_body_dict, BODY_IDENTIFIERS
from spoter_mod.normalization.hand_normalization import normalize_single_dict as normalize_single_hand_dict, HAND_IDENTIFIERS
model = torch.load("spoter-checkpoint.pth", map_location=torch.device('cpu'))
model.train(False)
HAND_IDENTIFIERS = [id + "_Left" for id in HAND_IDENTIFIERS] + [id + "_Right" for id in HAND_IDENTIFIERS]
GLOSS = ['book', 'drink', 'computer', 'before', 'chair', 'go', 'clothes', 'who', 'candy', 'cousin', 'deaf', 'fine',
'help', 'no', 'thin', 'walk', 'year', 'yes', 'all', 'black', 'cool', 'finish', 'hot', 'like', 'many', 'mother',
'now', 'orange', 'table', 'thanksgiving', 'what', 'woman', 'bed', 'blue', 'bowling', 'can', 'dog', 'family',
'fish', 'graduate', 'hat', 'hearing', 'kiss', 'language', 'later', 'man', 'shirt', 'study', 'tall', 'white',
'wrong', 'accident', 'apple', 'bird', 'change', 'color', 'corn', 'cow', 'dance', 'dark', 'doctor', 'eat',
'enjoy', 'forget', 'give', 'last', 'meet', 'pink', 'pizza', 'play', 'school', 'secretary', 'short', 'time',
'want', 'work', 'africa', 'basketball', 'birthday', 'brown', 'but', 'cheat', 'city', 'cook', 'decide', 'full',
'how', 'jacket', 'letter', 'medicine', 'need', 'paint', 'paper', 'pull', 'purple', 'right', 'same', 'son',
'tell', 'thursday']
device = torch.device("cpu")
if torch.cuda.is_available():
device = torch.device("cuda")
def tensor_to_dictionary(landmarks_tensor: torch.Tensor) -> dict:
data_array = landmarks_tensor.numpy()
output = {}
for landmark_index, identifier in enumerate(BODY_IDENTIFIERS + HAND_IDENTIFIERS):
output[identifier] = data_array[:, landmark_index]
return output
def dictionary_to_tensor(landmarks_dict: dict) -> torch.Tensor:
output = np.empty(shape=(len(landmarks_dict["leftEar"]), len(BODY_IDENTIFIERS + HAND_IDENTIFIERS), 2))
for landmark_index, identifier in enumerate(BODY_IDENTIFIERS + HAND_IDENTIFIERS):
output[:, landmark_index, 0] = [frame[0] for frame in landmarks_dict[identifier]]
output[:, landmark_index, 1] = [frame[1] for frame in landmarks_dict[identifier]]
return torch.from_numpy(output)
def greet(video):
data = obtain_pose_data(video)
depth_map = np.empty(shape=(len(data.data_hub["nose_X"]), len(BODY_IDENTIFIERS + HAND_IDENTIFIERS), 2))
for index, identifier in enumerate(BODY_IDENTIFIERS + HAND_IDENTIFIERS):
depth_map[:, index, 0] = data.data_hub[identifier + "_X"]
depth_map[:, index, 1] = data.data_hub[identifier + "_Y"]
depth_map = torch.from_numpy(np.copy(depth_map))
print(depth_map.shape)
depth_map = tensor_to_dictionary(depth_map)
keys = copy.copy(list(depth_map.keys()))
for key in keys:
data = depth_map[key]
del depth_map[key]
depth_map[key.replace("_Left", "_0").replace("_Right", "_1")] = data
depth_map = normalize_single_body_dict(depth_map)
depth_map = normalize_single_hand_dict(depth_map)
keys = copy.copy(list(depth_map.keys()))
for key in keys:
data = depth_map[key]
del depth_map[key]
depth_map[key.replace("_0", "_Left").replace("_1", "_Right")] = data
depth_map = dictionary_to_tensor(depth_map)
depth_map = depth_map - 0.5
inputs = depth_map.squeeze(0).to(device)
outputs = model(inputs).expand(1, -1, -1)
results = torch.nn.functional.softmax(outputs, dim=2).detach().numpy()[0, 0]
results = {GLOSS[i]: float(results[i]) for i in range(100)}
print(outputs.shape)
print(results)
return results
label = gr.outputs.Label(num_top_classes=5)
demo = gr.Interface(fn=greet, inputs="playable_video", outputs=label)
demo.launch()
|