import streamlit as st  
import numpy as np  
import cv2  
import tempfile  
from gradio_client import Client, handle_file  
from PIL import Image  

# Проверка доступности нового API  
client = Client("yeecin/img2text")  

# Заголовок приложения  
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:  
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  
            pil_image = Image.fromarray(frame_rgb)  
            st.image(pil_image, caption=f"Random Frame {random_frame}")  

            # Сохранение кадра во временный файл для API  
            buf = tempfile.NamedTemporaryFile(suffix='.jpg', delete=False)  
            pil_image.save(buf, format='JPEG')  
            buf.close()  

            try:  
                # Вызов нового API для получения описания  
                result = client.predict(  
                    raw_image=handle_file(buf.name),  
                    model_n="Image Captioning",  
                    strategy="Nucleus sampling",  
                    api_name="/predict"  
                )  
                description = result  
                st.success(f"Generated Description: {description}")  
            except Exception as e:  
                st.error(f"Error: Could not get a response from the model. {str(e)}")  
        else:  
            st.error("Error: Could not read a frame from the video.")  
    else:  
        st.error("Error: Video file does not contain any frames.")  

if cap is not None:  
    cap.release()