File size: 2,472 Bytes
4a620fb |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import gradio as gr
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from PIL import Image
import cv2
import torch
# Load model and processor
mix_model_id = "google/paligemma-3b-mix-224"
mix_model = PaliGemmaForConditionalGeneration.from_pretrained(mix_model_id)
mix_processor = AutoProcessor.from_pretrained(mix_model_id)
# Define function to extract frames from the video
def extract_frames(video_path, frame_interval=1):
# Open the video file
vidcap = cv2.VideoCapture(video_path)
frames = []
success, image = vidcap.read()
count = 0
while success:
# Capture a frame at the specified interval
if count % frame_interval == 0:
frames.append(image)
success, image = vidcap.read()
count += 1
vidcap.release()
return frames
# Define function to generate captions for a video
def process_video(video, prompt):
# Use video directly as the path (video is passed as a string)
frames = extract_frames(video, frame_interval=10) # Extract frames at intervals
captions = []
for frame in frames:
# Convert frame to PIL Image and process it (assuming mix_processor handles PIL Image)
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
inputs = mix_processor(image.convert("RGB"), prompt, return_tensors="pt")
try:
# Generate output from the model for each frame
output = mix_model.generate(**inputs, max_new_tokens=20)
# Decode and store the output for the frame
decoded_output = mix_processor.decode(output[0], skip_special_tokens=True)
captions.append(decoded_output[len(prompt):]) # Remove prompt part from the output
except IndexError as e:
print(f"IndexError: {e}")
captions.append("Error processing frame")
# Combine all frame captions into a coherent video description
return " ".join(captions)
# Define Gradio interface for video captioning
inputs = [
gr.Video(label="Upload Video"),
gr.Textbox(label="Prompt", placeholder="Enter your question")
]
outputs = gr.Textbox(label="Generated Caption")
# Create the Gradio app for video captioning
demo = gr.Interface(fn=process_video, inputs=inputs, outputs=outputs, title="Video Captioning with Mix PaliGemma Model",
description="Upload a video and get captions based on your prompt.")
# Launch the app
demo.launch(debug=True) |