import streamlit as st
import numpy as np
import cv2
import torch
from PIL import Image
from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel

# Инициализация модели
model_id = "nttdataspain/vit-gpt2-stablediffusion2-lora"
model = VisionEncoderDecoderModel.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
feature_extractor = ViTFeatureExtractor.from_pretrained(model_id)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Функция для получения текста из изображения
def predict(image):
    img = image.convert('RGB')
    model.eval()
    pixel_values = feature_extractor(images=[img], return_tensors="pt").pixel_values.to(device)
    with torch.no_grad():
        output_ids = model.generate(pixel_values, max_length=16, num_beams=4, return_dict_in_generate=True).sequences

    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    return preds[0]

# Streamlit интерфейс
st.title("Video Frame to Image Description")

# Загрузка видеофайла
uploaded_file = st.file_uploader("Upload a video file", type=["mp4", "avi", "mov"])

cap = None  # Инициализируем объект cap как None

if uploaded_file is not None:
    # Создаем временный файл для хранения видео
    tfile = tempfile.NamedTemporaryFile(delete=False)
    tfile.write(uploaded_file.read())
    
    # Захват видео
    cap = cv2.VideoCapture(tfile.name)
    length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    if length > 0:
        # Выбор случайного кадра
        random_frame = np.random.randint(length)
        cap.set(cv2.CAP_PROP_POS_FRAMES, random_frame)
        ret, frame = cap.read()

        if ret:
            # Конвертация кадра OpenCV в PIL Image
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(frame_rgb)

            # Отображение выбранного кадра
            st.image(pil_image, caption=f"Random Frame {random_frame}")

            # Получение текста из изображения
            description = predict(pil_image)
            st.success(f"Generated Description: {description}")
        else:
            st.error("Error: Could not read a frame from the video.")
    else:
        st.error("Error: Video file does not contain any frames.")

# Проверяем, был ли cap создан, и только тогда освобождаем ресурсы
if cap is not None:
    cap.release()