47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Part of the implementation is borrowed and modified from DUTCode,
|
|
# publicly available at https://github.com/Annbless/DUTCode
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from modelscope.preprocessors.cv import VideoReader
|
|
|
|
|
|
def stabilization_preprocessor(input, cfg):
|
|
video_reader = VideoReader(input)
|
|
inputs = []
|
|
for frame in video_reader:
|
|
inputs.append(np.flip(frame, axis=2))
|
|
fps = video_reader.fps
|
|
w = video_reader.width
|
|
h = video_reader.height
|
|
rgb_images = []
|
|
images = []
|
|
ori_images = []
|
|
for i, frame in enumerate(inputs):
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
|
image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
image = image * (1. / 255.)
|
|
image = cv2.resize(image, (cfg.MODEL.WIDTH, cfg.MODEL.HEIGHT))
|
|
images.append(image.reshape(1, 1, cfg.MODEL.HEIGHT, cfg.MODEL.WIDTH))
|
|
rgb_image = cv2.resize(frame, (cfg.MODEL.WIDTH, cfg.MODEL.HEIGHT))
|
|
rgb_images.append(
|
|
np.expand_dims(np.transpose(rgb_image, (2, 0, 1)), 0))
|
|
ori_images.append(np.expand_dims(np.transpose(frame, (2, 0, 1)), 0))
|
|
x = np.concatenate(images, 1).astype(np.float32)
|
|
x = torch.from_numpy(x).unsqueeze(0)
|
|
x_rgb = np.concatenate(rgb_images, 0).astype(np.float32)
|
|
x_rgb = torch.from_numpy(x_rgb).unsqueeze(0)
|
|
|
|
return {
|
|
'ori_images': ori_images,
|
|
'x': x,
|
|
'x_rgb': x_rgb,
|
|
'fps': fps,
|
|
'width': w,
|
|
'height': h
|
|
}
|