File size: 6,194 Bytes
f799d96 b0a48de f799d96 987b643 f799d96 987b643 f799d96 b0a48de f799d96 987b643 f799d96 b0a48de f799d96 b0a48de f799d96 |
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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
import gradio as gr
import cv2
import os
import tempfile
from torchvision import transforms
from torchvision.transforms import Compose
import torch
import numpy as np
from PIL import Image
import torch.nn.functional as F
from transformers import pipeline, TimesformerModel, VideoMAEImageProcessor
from utils import *
from algorithm import *
def make_video(video_path, outdir='./summarized_video',encoder='Kmeans'):
if encoder not in ["Kmeans", "Sum of Squared Difference 01", "Sum of Squared Difference 02"]:
encoder = "Kmeans"
# nen them vao cac truong hop mo hinh khac
margin_width = 50
model, processor, device = load_model()
# total_params = sum(param.numel() for param in model.parameters())
# print('Total parameters: {:.2f}M'.format(total_params / 1e6))
if os.path.isfile(video_path):
if video_path.endswith('txt'):
with open(video_path, 'r') as f:
lines = f.read().splitlines()
else:
filenames = [video_path]
else:
filenames = os.listdir(video_path)
filenames = [os.path.join(video_path, filename) for filename in filenames if not filename.startswith('.')]
filenames.sort()
for k, filename in enumerate(filenames):
print('Progress {:}/{:},'.format(k+1, len(filenames)), 'Processing', filename)
raw_video = cv2.VideoCapture(filename)
frame_width, frame_height = int(raw_video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(raw_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_rate = int(raw_video.get(cv2.CAP_PROP_FPS))
#length = int(raw_video.get(cv2.CAP_PROP_FRAME_COUNT))
output_width = frame_width * 2 + margin_width
filename = os.path.basename(filename)
# Find the size to resize
if "shortest_edge" in processor.size:
height = width = processor.size["shortest_edge"]
else:
height = processor.size["height"]
width = processor.size["width"]
resize_to = (height, width)
# F/Fs
clip_sample_rate = 1
# F
num_frames = 8
frames = []
features = []
# output_path = os.path.join(outdir, filename[:filename.rfind('.')] + '_video_depth.mp4')
with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmpfile:
output_path = tmpfile.name
#out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"avc1"), frame_rate, (output_width, frame_height))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, frame_rate, (output_width, frame_height))
# count=0
while raw_video.isOpened():
ret, raw_frame = raw_video.read()
if not ret:
break
raw_frame = cv2.resize(raw_frame, resize_to)
frames.append(raw_frame)
# Find key frames by selecting frames with clip_sample_rate
key_frames = frames[::clip_sample_rate]
#print('total of frames after sample:', len(selected_frames))
# Remove redundant frames to make the number of frames can be divided by num_frames
num_redudant_frames = len(key_frames) - (len(key_frames) % num_frames)
# Final key frames
final_key_frames = key_frames[:num_redudant_frames]
#print('total of frames after remove redundant frames:', len(selected_frames))
for i in range(0, len(final_key_frames), num_frames):
if i % num_frames*50 == 0:
print(f"Loading {i}/{len(final_key_frames)}")
# Input clip to the model
input_frames = final_key_frames[i:i+num_frames]
# Extract features
batch_features = extract_features(input_frames, device, model, processor)
# Convert to numpy array to decrease the memory usage
batch_features = np.array(batch_features.cpu().detach().numpy())
features.extend(batch_features)
number_of_clusters = round(len(features)*0.15)
selected_frames = []
if encoder == "Kmeans":
selected_frames = kmeans(features, number_of_clusters)
elif encoder == "Sum of Squared Difference 01":
selected_frames = tt01(features, 400)
else:
selected_frames = tt02(features, 400)
video_writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), frame_rate, (frames[0].shape[1], frames[0].shape[0]))
for idx in selected_frames:
video_writer.write(frames[idx])
raw_video.release()
video_writer.release()
print("Completed summarizing the video (wait for a moment to load).")
return output_path
css = """
#img-display-container {
max-height: 100vh;
}
#img-display-input {
max-height: 80vh;
}
#img-display-output {
max-height: 80vh;
}
"""
title = "# Video Summarization Demo"
description = """Video Summarization using Timesformer.
Author: Nguyen Hoai Nam.
"""
with gr.Blocks(css=css) as demo:
gr.Markdown(title)
gr.Markdown(description)
gr.Markdown("### Video Summarization demo")
with gr.Row():
input_video = gr.Video(label="Input Video")
algorithm_type = gr.Dropdown(["Kmeans", "Sum of Squared Difference 01", "Sum of Squared Difference 02"], type="value", label='Algorithm')
submit = gr.Button("Submit")
processed_video = gr.Video(label="Summarized Video")
def on_submit(uploaded_video,algorithm_type):
# Process the video and get the path of the output video
#output_video_path = make_video(uploaded_video,encoder=model_type)
pass
#return output_video_path
submit.click(on_submit, inputs=[input_video, algorithm_type], outputs=processed_video)
#example_files = os.listdir('assets/examples_video')
#example_files.sort()
#example_files = [os.path.join('assets/examples_video', filename) for filename in example_files]
#examples = gr.Examples(examples=example_files, inputs=[input_video], outputs=processed_video, fn=on_submit, cache_examples=True)
if __name__ == '__main__':
demo.queue().launch() |