app.py

import os
import cv2
import torch
import numpy as np
import threading
import json
import shutil 
from flask import Flask, request, jsonify, send_from_directory, render_template
from flask_cors import CORS
from werkzeug.utils import secure_filename
from ultralytics import YOLO
from torch.utils.data import Dataset, DataLoader

# --- CẤU HÌNH ---
UPLOAD_FOLDER = 'uploads'
DATA_FOLDER = 'data'
ALLOWED_EXTENSIONS = {'mp4', 'avi', 'mov'}
MODEL_SAVE_PATH = 'stgcn_model.pth'
CLASS_MAP_PATH = 'class_map.json'

app = Flask(__name__)
CORS(app) 
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['DATA_FOLDER'] = DATA_FOLDER

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(DATA_FOLDER, exist_ok=True)

# --- BIẾN TOÀN CỤC CHO VIỆC DỰ ĐOÁN ---
training_status = {'is_training': False, 'progress': 'Idle'}
prediction_model = None
prediction_class_map = None
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- ĐỊNH NGHĨA MÔ HÌNH ST-GCN (Không thay đổi) ---
class Graph:
    def __init__(self, layout='coco_pose', strategy='spatial'):
        self.num_node = 17
        self.self_link = [(i, i) for i in range(self.num_node)]
        self.inward = [(15, 13), (13, 11), (16, 14), (14, 12), (11, 12), (5, 11), (6, 12), (5, 6), (5, 7), (6, 8), (7, 9), (8, 10), (1, 2), (0, 1), (0, 2), (1, 3), (2, 4)]
        self.outward = [(j, i) for (i, j) in self.inward]
        self.neighbor = self.inward + self.outward
        self.A = self.get_adjacency_matrix(strategy)
    def get_adjacency_matrix(self, strategy):
        adjacency = np.zeros((self.num_node, self.num_node))
        for i, j in self.neighbor: adjacency[i, j] = 1; adjacency[j, i] = 1
        for i in range(self.num_node): adjacency[i, i] = 1
        D = np.array(adjacency.sum(1)); D_inv = np.power(D, -1).flatten(); D_inv[np.isinf(D_inv)] = 0.; D_mat_inv = np.diag(D_inv)
        norm_adj = D_mat_inv.dot(adjacency); A = np.zeros((1, self.num_node, self.num_node)); A[0, :, :] = norm_adj
        return A

class ST_GCN_Block(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dropout=0, residual=True):
        super().__init__(); assert len(kernel_size) == 2 and kernel_size[0] % 2 == 1; padding = ((kernel_size[0] - 1) // 2, 0)
        self.gcn = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1)
        self.tcn = torch.nn.Sequential(torch.nn.BatchNorm2d(out_channels), torch.nn.ReLU(inplace=True), torch.nn.Conv2d(out_channels, out_channels, (kernel_size[0], 1), (stride, 1), padding), torch.nn.BatchNorm2d(out_channels), torch.nn.Dropout(dropout, inplace=True))
        if not residual: self.residual = lambda x: 0
        elif (in_channels == out_channels) and (stride == 1): self.residual = lambda x: x
        else: self.residual = torch.nn.Sequential(torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=(stride, 1)), torch.nn.BatchNorm2d(out_channels))
        self.relu = torch.nn.ReLU(inplace=True)
    def forward(self, x, A):
        res = self.residual(x); x_gcn = self.gcn(x); x = torch.einsum('nctv,kvw->nctw', (x_gcn, A)); x = self.tcn(x); x = x + res
        return self.relu(x)

class STGCN(torch.nn.Module):
    def __init__(self, in_channels, num_classes, graph_args, edge_importance_weighting=True):
        super().__init__(); self.graph = Graph(**graph_args); A = torch.tensor(self.graph.A, dtype=torch.float32, requires_grad=False); self.register_buffer('A', A)
        spatial_kernel_size = A.size(0); temporal_kernel_size = 9; kernel_size = (temporal_kernel_size, spatial_kernel_size)
        self.data_bn = torch.nn.BatchNorm1d(in_channels * self.graph.num_node)
        self.st_gcn_networks = torch.nn.ModuleList((ST_GCN_Block(in_channels, 64, kernel_size, stride=1, residual=False),
                                                     ST_GCN_Block(64, 64, kernel_size, stride=1), 
                                                     ST_GCN_Block(64, 128, kernel_size, stride=2), 
                                                     ST_GCN_Block(128, 128, kernel_size, stride=1),
                                                    ST_GCN_Block(128, 256, kernel_size, stride=2),
                                                     ST_GCN_Block(256, 256, kernel_size, stride=1)))
        if edge_importance_weighting: self.edge_importance = torch.nn.ParameterList([torch.nn.Parameter(torch.ones(self.A.size())) for _ in self.st_gcn_networks])
        else: self.edge_importance = [1] * len(self.st_gcn_networks)
        self.fcn = torch.nn.Conv2d(256, num_classes, kernel_size=1)
    def forward(self, x):
        N, C, T, V = x.shape; x = x.permute(0, 3, 1, 2).contiguous(); x = x.view(N, V * C, T); x = self.data_bn(x); x = x.view(N, V, C, T).permute(0, 2, 3, 1).contiguous()
        for gcn, importance in zip(self.st_gcn_networks, self.edge_importance): importance_device = importance.to(x.device); A_weighted = self.A.to(x.device) * importance_device; x = gcn(x, A_weighted)
        x = torch.nn.functional.avg_pool2d(x, x.size()[2:]); x = x.view(N, -1, 1, 1); x = self.fcn(x); x = x.view(N, -1)
        return x

# --- CÁC HÀM HỖ TRỢ ---
def extract_pose_from_video(video_path, output_path, target_frames=200):
    print(f"Bắt đầu trích xuất pose từ video: {video_path} với target_frames={target_frames}")
    model = YOLO('yolo11n-pose.pt'); cap = cv2.VideoCapture(video_path); all_keypoints = []
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret: break
        results = model(frame, verbose=False)
        if results[0].keypoints and len(results[0].keypoints.xy[0]) > 0: all_keypoints.append(results[0].keypoints.xy[0].cpu().numpy())
        else: all_keypoints.append(np.zeros((17, 2)))
    cap.release()
    if not all_keypoints: print("Không có keypoints nào được trích xuất."); return False
    keypoints_sequence = np.array(all_keypoints); num_frames = keypoints_sequence.shape[0]
    if num_frames < target_frames: padding = np.tile(keypoints_sequence[-1:], (target_frames - num_frames, 1, 1)); keypoints_sequence = np.concatenate((keypoints_sequence, padding), axis=0)
    elif num_frames > target_frames: keypoints_sequence = keypoints_sequence[:target_frames]
    keypoints_sequence = np.transpose(keypoints_sequence, (2, 0, 1)); np.save(output_path, keypoints_sequence)
    print(f"Đã lưu dữ liệu pose vào: {output_path}"); return True

class PoseDataset(Dataset):
    def __init__(self, data_folder):
        self.samples = []; self.class_map = {}; class_idx = 0
        for class_name in sorted(os.listdir(data_folder)):
            class_dir = os.path.join(data_folder, class_name)
            if os.path.isdir(class_dir):
                if class_name not in self.class_map: self.class_map[class_name] = class_idx; class_idx += 1
                for sample_file in os.listdir(class_dir):
                    if sample_file.endswith('.npy'): self.samples.append((os.path.join(class_dir, sample_file), self.class_map[class_name]))
    def __len__(self): return len(self.samples)
    def __getitem__(self, idx):
        sample_path, label = self.samples[idx]; data = np.load(sample_path)
        return torch.from_numpy(data).float(), torch.tensor(label).long()

def train_model_thread(num_epochs=50):
    global training_status, prediction_model, prediction_class_map
    training_status['is_training'] = True; training_status['progress'] = 'Starting training...'
    try:
        dataset = PoseDataset(app.config['DATA_FOLDER'])
        if len(dataset) == 0: raise ValueError("No data found. Please upload videos first.")
        num_classes = len(dataset.class_map)
        if num_classes < 2: raise ValueError("Need at least 2 classes to train.")
        
        with open(CLASS_MAP_PATH, 'w') as f:
            json.dump(dataset.class_map, f)
        print(f"Class map đã được lưu tại {CLASS_MAP_PATH}")

        print(f"Bắt đầu training với {len(dataset)} mẫu, {num_classes} lớp, trong {num_epochs} epochs."); print("Lớp được ánh xạ:", dataset.class_map)
        data_loader = DataLoader(dataset, batch_size=16, shuffle=True)
        model = STGCN(in_channels=2, num_classes=num_classes, graph_args={}).to(device)
        optimizer = torch.optim.Adam(model.parameters(), lr=0.005); criterion = torch.nn.CrossEntropyLoss()
        for epoch in range(num_epochs):
            model.train(); running_loss = 0.0
            for i, (inputs, labels) in enumerate(data_loader):
                inputs, labels = inputs.to(device), labels.to(device); optimizer.zero_grad(); outputs = model(inputs)
                loss = criterion(outputs, labels); loss.backward(); optimizer.step(); running_loss += loss.item()
            epoch_loss = running_loss / len(data_loader); progress_text = f"Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}"
            print(progress_text); training_status['progress'] = progress_text
        torch.save(model.state_dict(), MODEL_SAVE_PATH); print(f"Training hoàn tất. Model đã được lưu tại {MODEL_SAVE_PATH}")
        training_status['progress'] = f'Training complete! Model saved to {MODEL_SAVE_PATH}.'
        
        load_model_for_prediction()

    except Exception as e: error_msg = f"An error occurred during training: {e}"; print(error_msg); training_status['progress'] = error_msg
    finally: training_status['is_training'] = False

def load_model_for_prediction():
    global prediction_model, prediction_class_map
    if os.path.exists(MODEL_SAVE_PATH) and os.path.exists(CLASS_MAP_PATH):
        try:
            with open(CLASS_MAP_PATH, 'r') as f: class_map = json.load(f)
            num_classes = len(class_map)
            model = STGCN(in_channels=2, num_classes=num_classes, graph_args={})
            model.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location=device)); model.to(device); model.eval()
            prediction_model = model; prediction_class_map = class_map
            print(f"Model '{MODEL_SAVE_PATH}' và class map đã được tải thành công.")
        except Exception as e: print(f"Lỗi khi tải model: {e}")
    else: print("Không tìm thấy model hoặc class map. Vui lòng training trước.")

def clear_data_directory():
    """Xóa tất cả dữ liệu trong thư mục data và các file model đã lưu."""
    if os.path.isdir(DATA_FOLDER):
        try:
            shutil.rmtree(DATA_FOLDER)
            print(f"Đã xóa thành công thư mục '{DATA_FOLDER}'.")
        except OSError as e:
            print(f"Lỗi khi xóa thư mục {DATA_FOLDER}: {e.strerror}")
    
    os.makedirs(DATA_FOLDER, exist_ok=True)
    print(f"Đã tạo lại thư mục rỗng '{DATA_FOLDER}'.")
    
    if os.path.exists(MODEL_SAVE_PATH):
        os.remove(MODEL_SAVE_PATH)
        print(f"Đã xóa file model: {MODEL_SAVE_PATH}")
    if os.path.exists(CLASS_MAP_PATH):
        os.remove(CLASS_MAP_PATH)
        print(f"Đã xóa file class map: {CLASS_MAP_PATH}")

# --- FLASK API ENDPOINTS ---
@app.route('/')
def index():
    """Route để hiển thị trang chủ giới thiệu."""
    return render_template('index.html')


@app.route('/demo')
def index_logic():
    return  render_template( 'index_logic.html')


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/upload', methods=['POST'])
def upload_video():
    if 'video' not in request.files: return jsonify({'error': 'No video file part'}), 400
    file = request.files['video']; class_name = request.form['class_name']; target_frames = int(request.form.get('target_frames', 200))
    if file.filename == '': return jsonify({'error': 'No selected file'}), 400
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename); video_path = os.path.join(app.config['UPLOAD_FOLDER'], filename); file.save(video_path)
        class_data_folder = os.path.join(app.config['DATA_FOLDER'], class_name); os.makedirs(class_data_folder, exist_ok=True)
        output_npy_path = os.path.join(class_data_folder, os.path.splitext(filename)[0] + '.npy')
        success = extract_pose_from_video(video_path, output_npy_path, target_frames=target_frames); os.remove(video_path)
        if success: return jsonify({'message': f'Video processed for class "{class_name}"'}), 200
        else: return jsonify({'error': 'Failed to extract pose from video'}), 500
    return jsonify({'error': 'File type not allowed'}), 400

@app.route('/train', methods=['POST'])
def trigger_training():
    if training_status['is_training']: return jsonify({'message': 'A training process is already running.'}), 409
    data = request.get_json(); num_epochs = data.get('epochs', 50)
    train_thread = threading.Thread(target=train_model_thread, kwargs={'num_epochs': num_epochs}); train_thread.start()
    return jsonify({'message': f'Training started for {num_epochs} epochs.'}), 202

@app.route('/status', methods=['GET'])
def get_training_status(): return jsonify(training_status)

@app.route('/download', methods=['GET'])
def download_model():
    directory = os.getcwd()
    try: return send_from_directory(directory, MODEL_SAVE_PATH, as_attachment=True)
    except FileNotFoundError: return jsonify({"error": "Model file not found."}), 404

@app.route('/predict', methods=['POST'])
def predict():
    if prediction_model is None or prediction_class_map is None: return jsonify({'error': 'Model is not loaded. Please train a model first.'}), 503
    data = request.get_json()
    if 'keypoints' not in data: return jsonify({'error': 'Missing keypoints data.'}), 400
    try:
        keypoints = np.array(data['keypoints']); keypoints = np.transpose(keypoints, (2, 0, 1)); keypoints = np.expand_dims(keypoints, axis=0)
        input_tensor = torch.from_numpy(keypoints).float().to(device)
        with torch.no_grad():
            output = prediction_model(input_tensor)
            probabilities = torch.nn.functional.softmax(output, dim=1).squeeze().cpu().numpy()
        idx_to_class = {v: k for k, v in prediction_class_map.items()}
        results = [{'className': idx_to_class.get(i, f"Unknown"), 'probability': float(prob)} for i, prob in enumerate(probabilities)]
        results.sort(key=lambda x: x['probability'], reverse=True)
        return jsonify(results)
    except Exception as e: print(f"Prediction error: {e}"); return jsonify({'error': 'An error occurred during prediction.'}), 500

# <<-- ENDPOINT MỚI ĐỂ EXPORT FILE ZIP -->>
@app.route('/export', methods=['GET'])
def export_zip():
    """Phục vụ file test.zip để tải về."""
    directory = os.getcwd()
    zip_filename = 'test.zip'
    if not os.path.exists(os.path.join(directory, zip_filename)):
        return jsonify({"error": f"File '{zip_filename}' not found on the server. Please add it to the project folder."}), 404
    try:
        return send_from_directory(directory, zip_filename, as_attachment=True)
    except FileNotFoundError:
        return jsonify({"error": "File not found during send operation."}), 404

# --- CHẠY APP ---
if __name__ == '__main__':
    clear_data_directory() # <<-- XÓA DỮ LIỆU CŨ KHI KHỞI ĐỘNG
    load_model_for_prediction()
    app.run(host='0.0.0.0', port=5000, debug=True)