Views
No views yet
1import os
2os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
3
4
5from transformers import VideoMAEImageProcessor, AutoModelForVideoClassification
6import torch
7import numpy as np
8import av
9
10
11
12def extract_frames(video_path, num_frames=16):
13 frames = []
14 try:
15 container = av.open(video_path)
16 stream = container.streams.video[0]
17 indices = np.linspace(0, stream.frames - 1, num=num_frames, dtype=int)
18
19 for i, frame in enumerate(container.decode(stream)):
20 if i in indices:
21 img = frame.to_ndarray(format="rgb24")
22 frames.append(img)
23 if len(frames) == num_frames:
24 break
25
26 container.close()
27 except Exception as e:
28 print(f"Erro ao processar o vídeo: {e}")
29
30 if len(frames) < num_frames:
31 pad_size = num_frames - len(frames)
32 if frames:
33 last_frame = frames[-1]
34 frames += [last_frame.copy() for _ in range(pad_size)]
35 else:
36 frames = [np.zeros((224, 224, 3), dtype=np.uint8)] * num_frames
37
38 return frames
39
40processor = VideoMAEImageProcessor.from_pretrained("maike616/pain-classifier-video")
41model = AutoModelForVideoClassification.from_pretrained("maike616/pain-classifier-video")
42
43high = r'AI4Pain Dataset\Validation\video\high_pain\1_Pain_HIGH_20.mp4'
44low = r'AI4Pain Dataset\Validation\video\low_pain\1_Pain_LOW_3.mp4'
45no_pain = r'AI4Pain Dataset\Validation\video\no_pain\1_Rest_1.mp4'
46
47video_frames = extract_frames(no_pain)
48
49inputs = processor(video_frames, return_tensors="pt")
50
51with torch.no_grad():
52 outputs = model(**inputs)
53 probs = torch.softmax(outputs.logits, dim=1)
54 predicted_class_idx = outputs.logits.argmax(-1).item()
55 predicted_class = model.config.id2label[predicted_class_idx]
56
57print(f"Classe predita: {predicted_class}")
58print(f"Probabilidades: {probs[0].tolist()}")
59