import torch
from transformers import AutoModelForImageClassification, AutoFeatureExtractor

class DrowsinessModel:
    def __init__(self, model_path, threshold=0.15):
        # Инициализация модели и экстрактора признаков
        self.model = AutoModelForImageClassification.from_pretrained('MonikaG7/fine-tuned_fatique_model')
        self.feature_extractor = AutoFeatureExtractor.from_pretrained('MonikaG7/fine-tuned_fatique_model')
        self.threshold = threshold

    def predict(self, video_path):
        """
        Анализирует видео, выводя 1 (усталость) или 0 (нет усталости)
        """
        import cv2
        cap = cv2.VideoCapture(video_path)
        predictions = []

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            inputs = self.feature_extractor(images=[frame], return_tensors="pt")
            with torch.no_grad():
                outputs = self.model(**inputs)
                pred = torch.argmax(outputs.logits, dim=1).item()
                predictions.append(pred)

        cap.release()

        if not predictions:
            return 0  # если видео пустое, возвращаем 0

        # Рассчитываем drowsiness_ratio
        drowsiness_ratio = sum(predictions) / len(predictions)
        return 1 if drowsiness_ratio > self.threshold else 0