| 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 |
