File size: 1,466 Bytes
5e40c7b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 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 |