File size: 2,614 Bytes
1239b39 |
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 |
import os
import cv2
import imageio
import numpy as np
def video_to_frame(video_path: str,
frame_dir: str,
filename_pattern: str = 'frame%03d.jpg',
log: bool = True,
frame_edit_func=None):
os.makedirs(frame_dir, exist_ok=True)
vidcap = cv2.VideoCapture(video_path)
success, image = vidcap.read()
if log:
print('img shape: ', image.shape[0:2])
count = 0
while success:
if frame_edit_func is not None:
image = frame_edit_func(image)
cv2.imwrite(os.path.join(frame_dir, filename_pattern % count), image)
success, image = vidcap.read()
if log:
print('Read a new frame: ', success, count)
count += 1
vidcap.release()
def frame_to_video(video_path: str, frame_dir: str, fps=30, log=True):
first_img = True
writer = imageio.get_writer(video_path, fps=fps)
file_list = sorted(os.listdir(frame_dir))
for file_name in file_list:
if not (file_name.endswith('jpg') or file_name.endswith('png')):
continue
fn = os.path.join(frame_dir, file_name)
curImg = imageio.imread(fn)
if first_img:
H, W = curImg.shape[0:2]
if log:
print('img shape', (H, W))
first_img = False
writer.append_data(curImg)
writer.close()
def get_fps(video_path: str):
video = cv2.VideoCapture(video_path)
fps = video.get(cv2.CAP_PROP_FPS)
video.release()
return fps
def get_frame_count(video_path: str):
video = cv2.VideoCapture(video_path)
frame_count = video.get(cv2.CAP_PROP_FRAME_COUNT)
video.release()
return frame_count
def resize_image(input_image, resolution):
H, W, C = input_image.shape
H = float(H)
W = float(W)
k = min(float(resolution) / min(H, W), float(768) / max(H, W))
H *= k
W *= k
H = int(np.round(H / 64.0)) * 64
W = int(np.round(W / 64.0)) * 64
img = cv2.resize(
input_image, (W, H),
interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
return img
def prepare_frames(input_path: str, output_dir: str, resolution: int, crop):
l, r, t, b = crop
def crop_func(frame):
H, W, C = frame.shape
left = np.clip(l, 0, W)
right = np.clip(W - r, left, W)
top = np.clip(t, 0, H)
bottom = np.clip(H - b, top, H)
frame = frame[top:bottom, left:right]
return resize_image(frame, resolution)
video_to_frame(input_path, output_dir, '%04d.png', False, crop_func)
|