SickstadiumAI / video_processing.py
jschwab21's picture
Update video_processing.py
91fc83c verified
raw
history blame
7.5 kB
import os
import cv2
from scenedetect import VideoManager, SceneManager
from scenedetect.detectors import ContentDetector
from moviepy.editor import VideoFileClip
from transformers import CLIPProcessor, CLIPModel
import torch
import yt_dlp
from PIL import Image
import uuid
from torchvision import models, transforms
from torch.nn import functional as F
import numpy as np
categories = ["Joy", "Trust", "Fear", "Surprise", "Sadness", "Disgust", "Anger", "Anticipation"]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
# Load ResNet-50 model
resnet50 = models.resnet50(pretrained=True)
resnet50.eval().to(device)
def classify_frame(frame):
# Preprocess the image
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(Image.fromarray(frame))
input_batch = input_tensor.unsqueeze(0).to(device)
# Predict with ResNet-50
with torch.no_grad():
output = resnet50(input_batch)
probabilities = F.softmax(output[0], dim=0)
# Create a numpy array from the probabilities of the categories
# This example assumes each category is mapped to a model output directly
results_array = np.array([probabilities[i].item() for i in range(len(categories))])
return results_array
def download_video(url):
ydl_opts = {
'format': 'bestvideo[height<=1440]+bestaudio/best[height<=1440]',
'outtmpl': f'temp_videos/{uuid.uuid4()}_video.%(ext)s',
'merge_output_format': 'mp4',
}
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
result = ydl.extract_info(url, download=True)
video_filename = ydl.prepare_filename(result)
safe_filename = sanitize_filename(video_filename)
if os.path.exists(video_filename) and video_filename != safe_filename:
os.rename(video_filename, safe_filename)
return safe_filename
def sanitize_filename(filename):
return "".join([c if c.isalnum() or c in " .-_()" else "_" for c in filename])
def find_scenes(video_path):
video_manager = VideoManager([video_path])
scene_manager = SceneManager()
scene_manager.add_detector(ContentDetector(threshold=33)) # Adjusted threshold for finer segmentation
video_manager.set_downscale_factor()
video_manager.start()
scene_manager.detect_scenes(frame_source=video_manager)
scene_list = scene_manager.get_scene_list()
video_manager.release()
scenes = [(start.get_timecode(), end.get_timecode()) for start, end in scene_list]
return scenes
def convert_timestamp_to_seconds(timestamp):
h, m, s = map(float, timestamp.split(':'))
return int(h) * 3600 + int(m) * 60 + s
def extract_frames(video, start_time, end_time):
frames = []
start_seconds = convert_timestamp_to_seconds(start_time)
end_seconds = convert_timestamp_to_seconds(end_time)
video_clip = video.subclip(start_seconds, end_seconds)
for frame_time in range(0, int(video_clip.duration * video_clip.fps), int(video_clip.fps / 4)):
frame = video_clip.get_frame(frame_time / video_clip.fps)
frames.append(frame)
return frames
def analyze_scenes(video_path, scenes, description, batch_size=65):
scene_scores = []
negative_descriptions = [
"black screen",
"Intro text for a video",
"dark scene without much contrast",
"No people are in this scene",
"A still shot of natural scenery",
"Still-camera shot of a person's face",
"One lazy dog on a log"
]
preprocess = transforms.Compose([
transforms.ToTensor(),
transforms.Resize((224, 224)),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
text_inputs = processor(text=[description] + negative_descriptions, return_tensors="pt", padding=True).to(device)
text_features = model.get_text_features(**text_inputs).detach()
positive_feature, negative_features = text_features[0], text_features[1:]
#print("Negative features shape:", negative_features.shape)
video = VideoFileClip(video_path)
for scene_num, (start_time, end_time) in enumerate(scenes):
frames = extract_frames(video, start_time, end_time)
if not frames:
print(f"Scene {scene_num + 1}: Start={start_time}, End={end_time} - No frames extracted")
continue
batches = [frames[i:i + batch_size] for i in range(0, len(frames), batch_size)]
scene_prob = 0.0
sentiment_distributions = np.zeros(8)
for batch in batches:
batch_tensors = torch.stack([preprocess(frame) for frame in batch]).to(device)
with torch.no_grad():
image_features = model.get_image_features(pixel_values=batch_tensors).detach()
#print("Image Features Shape:", image_features.shape)
positive_similarities = torch.cosine_similarity(image_features, positive_feature.unsqueeze(0).expand_as(image_features))
negative_mean = negative_features.mean(dim=0).unsqueeze(0).expand_as(image_features)
negative_similarities = torch.cosine_similarity(image_features, negative_mean)
scene_prob += (positive_similarities.mean().item() - negative_similarities.mean().item())
for frame in batch:
frame_sentiments = classify_frame(frame)
sentiment_distributions += np.array(frame_sentiments)
sentiment_distributions /= len(frames)
sentiment_percentages = {category: round(prob * 100, 2) for category, prob in zip(categories, sentiment_distributions)}
scene_prob /= len(frames)
scene_duration = convert_timestamp_to_seconds(end_time) - convert_timestamp_to_seconds(start_time)
print(f"Scene {scene_num + 1}: Start={start_time}, End={end_time}, Probability={scene_prob}, Duration={scene_duration}, Sentiments: {sentiment_percentages}")
scene_scores.append((scene_prob, start_time, end_time, scene_duration, sentiment_percentages))
scene_scores.sort(reverse=True, key=lambda x: x[0])
top_3_scenes = scene_scores[:3]
best_scene = max(top_3_scenes, key=lambda x: x[3])
if best_scene:
print(f"Best Scene: Start={best_scene[1]}, End={best_scene[2]}, Probability={best_scene[0]}, Duration={best_scene[3]}, Sentiments: {best_scene[4]}")
return (best_scene[1], best_scene[2]), best_scene[4]
else:
print("No suitable scene found")
return None, {}
def extract_best_scene(video_path, scene):
if scene is None:
return None
start_time, end_time = scene
start_seconds = convert_timestamp_to_seconds(start_time)
end_seconds = convert_timestamp_to_seconds(end_time)
video_clip = VideoFileClip(video_path).subclip(start_seconds, end_seconds)
return video_clip
def cleanup_temp_files():
temp_dir = 'temp_videos'
if os.path.exists(temp_dir):
for file in os.listdir(temp_dir):
file_path = os.path.join(temp_dir, file)
try:
if os.path.isfile(file_path):
os.unlink(file_path)
except Exception as e:
print(f"Error: {e}")