Create video_model.py

video_model.py

@@ -0,0 +1,37 @@
+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