Spaces:
Runtime error
Runtime error
File size: 1,704 Bytes
ff715ca |
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 |
import cv2
import torch.nn.functional as F
import numpy as np
from src.utils import *
def insert_key(keys, ind):
for i, k in enumerate(keys):
if ind < k:
keys.insert(i, ind)
break
def get_maxinterv(keys):
maxinterv = 1
for i in range(len(keys)-1):
tmp = keys[i+1]-keys[i]
if tmp > maxinterv:
maxinterv = tmp
return maxinterv
def get_keyframe_ind(filename, lastframen = 1e10, mininterv = 5, maxinterv = 20, viz = False):
if maxinterv == mininterv:
return list(range(0,lastframen,mininterv))
video_cap = cv2.VideoCapture(filename)
n_frames = max(1, min(int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT)), lastframen))
err = [0]
preframe = None
for i in range(n_frames):
success, frame = video_cap.read()
if not success:
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
img = resize_image(frame, 512)
img = cv2.GaussianBlur(img, (9, 9), 0.0)
if i == 0:
preframe = numpy2tensor(img)
else:
curframe = numpy2tensor(img)
err += [float(F.mse_loss(preframe, curframe).cpu().numpy())]
preframe = curframe
err = np.array(err)
err1 = np.array(err)
n_frames = len(err)
keys = [0, n_frames-1]
err[0:mininterv] = -1
err[-mininterv:] = -1
while get_maxinterv(keys) > maxinterv:
ind = np.argmax(err)
if err[ind] == -1:
break
err[ind-mininterv:ind+mininterv] = -1
insert_key(keys, ind)
if viz:
plt.plot(err1)
plt.plot(keys, err1[keys], 'bo')
plt.show()
return keys |