246 lines
6.9 KiB
Python
246 lines
6.9 KiB
Python
"""
|
|
Inference utilities untuk real-time gesture recognition.
|
|
|
|
Melakukan prediksi menggunakan model TCN sesuai README.
|
|
"""
|
|
|
|
import cv2
|
|
import json
|
|
import mediapipe as mp
|
|
import numpy as np
|
|
import torch
|
|
from pathlib import Path
|
|
from typing import Optional, Tuple
|
|
|
|
from utils.mediapipe_detector import MediaPipeDetector
|
|
from utils.normalization import LandmarkNormalizer
|
|
from utils.config import WINDOW_SIZE
|
|
|
|
|
|
# MediaPipe Drawing utils
|
|
mp_drawing = mp.solutions.drawing_utils
|
|
mp_hands = mp.solutions.hands
|
|
mp_face_mesh = mp.solutions.face_mesh
|
|
|
|
|
|
class GestureRecognizer:
|
|
"""
|
|
Real-time gesture recognizer untuk inferensi.
|
|
|
|
Proses sesuai README:
|
|
1. Webcam
|
|
2. MediaPipe
|
|
3. Landmark
|
|
4. Normalisasi
|
|
5. Window 30 Frame
|
|
6. Model TCN
|
|
7. Prediksi
|
|
8. Tampilkan hasil
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_path: Path,
|
|
label_encoder_path: Path,
|
|
window_size: int = WINDOW_SIZE,
|
|
smoothing_size: int = 5,
|
|
draw_landmarks: bool = True
|
|
):
|
|
"""
|
|
Inisialisasi recognizer.
|
|
|
|
Args:
|
|
model_path: Path ke model .pt
|
|
label_encoder_path: Path ke label encoder JSON
|
|
window_size: Window size untuk sequence (30)
|
|
smoothing_size: Buffer size untuk temporal smoothing
|
|
draw_landmarks: Apakah perlu menggambar landmarks
|
|
"""
|
|
self.window_size = window_size
|
|
self.smoothing_size = smoothing_size
|
|
self.draw_landmarks = draw_landmarks
|
|
|
|
# Load model
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
self.model = torch.load(model_path, map_location=self.device, weights_only=False)
|
|
self.model.eval()
|
|
|
|
# Load label encoder
|
|
with open(label_encoder_path, 'r') as f:
|
|
encoder_data = json.load(f)
|
|
self.idx_to_label = encoder_data["idx_to_label"]
|
|
|
|
# Initialize detector dan normalizer
|
|
self.detector = MediaPipeDetector()
|
|
self.normalizer = LandmarkNormalizer()
|
|
|
|
# Buffer untuk temporal smoothing
|
|
self.buffer = []
|
|
self.prob_buffer = []
|
|
|
|
# MediaPipe drawing specs
|
|
self.hand_landmark_spec = mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=3)
|
|
self.hand_connection_spec = mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2)
|
|
self.face_landmark_spec = mp_drawing.DrawingSpec(color=(255, 255, 0), thickness=1, circle_radius=1)
|
|
|
|
def process_frame(self, frame) -> Tuple[Optional[str], float]:
|
|
"""
|
|
Process satu frame dan return prediction.
|
|
|
|
Args:
|
|
frame: Frame gambar (BGR)
|
|
|
|
Returns:
|
|
Tuple (predicted_label, confidence)
|
|
"""
|
|
# Extract landmarks
|
|
left_hand, right_hand, face = self.detector.detect(frame)
|
|
|
|
# Check jika ada gerakan (minimal satu tangan terdeteksi)
|
|
left_detected = np.any(left_hand != 0)
|
|
right_detected = np.any(right_hand != 0)
|
|
|
|
# Jika tidak ada tangan terdeteksi beberapa frame berturut-turut, reset buffer
|
|
if not (left_detected or right_detected):
|
|
self.no_hand_count = getattr(self, 'no_hand_count', 0) + 1
|
|
if self.no_hand_count >= 5:
|
|
self.buffer = []
|
|
self.prob_buffer = []
|
|
return None, 0.0
|
|
|
|
self.no_hand_count = 0
|
|
|
|
# Normalize
|
|
features = self.normalizer.normalize_frame(left_hand, right_hand, face)
|
|
|
|
# Add ke buffer
|
|
self.buffer.append(features)
|
|
|
|
# Keep buffer size
|
|
if len(self.buffer) > self.window_size:
|
|
self.buffer.pop(0)
|
|
|
|
# Jika buffer belum penuh, return None
|
|
if len(self.buffer) < self.window_size:
|
|
return None, 0.0
|
|
|
|
# Prepare input
|
|
sequence = np.array([self.buffer])
|
|
sequence_tensor = torch.FloatTensor(sequence).to(self.device)
|
|
|
|
# Predict
|
|
with torch.no_grad():
|
|
outputs = self.model(sequence_tensor)
|
|
probabilities = torch.softmax(outputs, dim=1)
|
|
|
|
# Store probabilities for smoothing
|
|
self.prob_buffer.append(probabilities.cpu().numpy()[0])
|
|
if len(self.prob_buffer) > self.smoothing_size:
|
|
self.prob_buffer.pop(0)
|
|
|
|
# Average probabilities
|
|
avg_probs = np.mean(self.prob_buffer, axis=0)
|
|
predicted_idx = np.argmax(avg_probs)
|
|
confidence = avg_probs[predicted_idx]
|
|
|
|
# Map ke label
|
|
label = self.idx_to_label[str(predicted_idx)]
|
|
|
|
return label, confidence
|
|
|
|
def draw_landmarks_on_frame(self, frame):
|
|
"""
|
|
Gambar hand dan face landmarks pada frame.
|
|
|
|
Args:
|
|
frame: Frame gambar (BGR)
|
|
|
|
Returns:
|
|
Frame dengan landmark yang sudah digambar
|
|
"""
|
|
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
# Detect hands
|
|
hands_results = self.detector.hands.process(rgb_frame)
|
|
|
|
# Detect face
|
|
face_results = self.detector.face_mesh.process(rgb_frame)
|
|
|
|
# Draw hands
|
|
if hands_results.multi_hand_landmarks:
|
|
for hand_landmarks in hands_results.multi_hand_landmarks:
|
|
mp_drawing.draw_landmarks(
|
|
frame,
|
|
hand_landmarks,
|
|
mp_hands.HAND_CONNECTIONS,
|
|
self.hand_landmark_spec,
|
|
self.hand_connection_spec
|
|
)
|
|
|
|
# Draw face mesh
|
|
if face_results.multi_face_landmarks:
|
|
for face_landmarks in face_results.multi_face_landmarks:
|
|
mp_drawing.draw_landmarks(
|
|
frame,
|
|
face_landmarks,
|
|
mp_face_mesh.FACEMESH_TESSELATION,
|
|
self.face_landmark_spec,
|
|
self.face_landmark_spec
|
|
)
|
|
|
|
return frame
|
|
|
|
def reset(self):
|
|
"""Reset buffers."""
|
|
self.buffer = []
|
|
self.prob_buffer = []
|
|
|
|
def close(self):
|
|
"""Cleanup resources."""
|
|
self.detector.close()
|
|
|
|
|
|
def load_model(model_path: Path, device: str = None):
|
|
"""
|
|
Load model dari file.
|
|
|
|
Args:
|
|
model_path: Path ke model file
|
|
device: Target device
|
|
|
|
Returns:
|
|
Loaded model
|
|
"""
|
|
if device is None:
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
model = torch.load(model_path, map_location=device, weights_only=False)
|
|
model.eval()
|
|
|
|
return model
|
|
|
|
|
|
def load_label_encoder(json_path: Path) -> dict:
|
|
"""
|
|
Load label encoder dari JSON file.
|
|
|
|
Args:
|
|
json_path: Path ke label encoder JSON
|
|
|
|
Returns:
|
|
Label encoder dictionary
|
|
"""
|
|
with open(json_path, 'r') as f:
|
|
return json.load(f)
|
|
|
|
|
|
def save_label_encoder(label_encoder: dict, json_path: Path):
|
|
"""
|
|
Save label encoder ke JSON file.
|
|
|
|
Args:
|
|
label_encoder: Label encoder dictionary
|
|
json_path: Output path
|
|
"""
|
|
with open(json_path, 'w') as f:
|
|
json.dump(label_encoder, f, indent=2) |