CristianLazoQuispe's picture
demo v1
979a0f3
import cv2
import torch
import numpy as np
import pandas as pd
from .islr.islr_model import DummyISLRModel
from .pose.keypoints_utils import extract_keypoints_from_frame
#model = DummyISLRModel(num_classes=100)
#model.load_state_dict(torch.load("demo_model.pt", map_location='cpu'))
#model.eval()
LABELS = [f"Clase {i}" for i in range(100)]
def predict_from_video(video_path,model=None):
cap = cv2.VideoCapture(video_path)
keypoints = []
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
keypoint = extract_keypoints_from_frame(frame)
keypoints.append(keypoint)
cap.release()
if not keypoints:
return "No keypoints detected", pd.DataFrame()
x = torch.tensor(np.mean(keypoints, axis=0)).float().unsqueeze(0)
#print("x:")
#print(x)
with torch.no_grad():
logits = model(x)
probs = torch.softmax(logits, dim=1).numpy()[0]
#print("probs:")
#print(probs)
top5_idx = probs.argsort()[-5:][::-1]
top5_labels = [LABELS[i] for i in top5_idx]
top5_probs = [probs[i] for i in top5_idx]
confidences = {LABELS[i]: float(probs[i]) for i in top5_idx}
#print("confidences:")
#print(confidences)
#df = pd.DataFrame({"label": top5_labels, "value": top5_probs})
return top5_labels[0],confidences # df