207 lines
6.1 KiB
Python
207 lines
6.1 KiB
Python
import argparse
|
|
import imageio
|
|
import os
|
|
import subprocess
|
|
|
|
|
|
def split_frame(videopath,
|
|
out_dir,
|
|
fps=30,
|
|
ext="jpg",
|
|
down_scale=1,
|
|
start_sec=0,
|
|
end_sec=-1,
|
|
overwrite=False,
|
|
**kwargs,):
|
|
import cv2
|
|
base_name = videopath.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
|
if os.path.isfile(videopath):
|
|
print(f'processing {videopath} .....')
|
|
# filepath = os.path.dirname(videopath)
|
|
# videoname = os.path.basename(videopath)[:-4]
|
|
if not os.path.exists(videopath):
|
|
raise('{} not exists!'.format(videopath))
|
|
|
|
vc = cv2.VideoCapture(videopath)
|
|
|
|
path = out_dir
|
|
if not os.path.exists(path):
|
|
os.makedirs(path)
|
|
|
|
if vc.isOpened():
|
|
rval , frame = vc.read()
|
|
if '.MOV' not in videopath:
|
|
cv2.imencode('.jpg',frame)[1].tofile(path +'/' + str(1).zfill(6) + '.jpg')
|
|
else:
|
|
cv2.imencode('.jpg',frame[::-1, ::-1])[1].tofile(path +'/' + str(0).zfill(6) + '.jpg')
|
|
else:
|
|
rval = False
|
|
|
|
Frame = 2
|
|
timeF = 1
|
|
while rval:
|
|
rval,frame = vc.read()
|
|
if rval==False:
|
|
break
|
|
|
|
if '.MOV' not in videopath:
|
|
cv2.imencode('.jpg',frame)[1].tofile(path +'/' + str(Frame).zfill(6) + '.jpg')
|
|
else:
|
|
cv2.imencode('.jpg',frame[::-1, ::-1])[1].tofile(path +'/' + str(Frame).zfill(6) + '.jpg')
|
|
|
|
Frame += 1
|
|
elif os.path.isdir(videopath):
|
|
print()
|
|
print('moving ', base_name , 'to ', out_dir)
|
|
os.system(f"cp -r {videopath} {out_dir}")
|
|
raise ValueError
|
|
return 0
|
|
|
|
else:
|
|
print(os.path.abspath(videopath), os.path.exists(videopath))
|
|
raise ValueError(f"{videopath} does not exist or is not a valid file or directory.")
|
|
|
|
return 0
|
|
|
|
def video_to_frames(
|
|
path,
|
|
out_dir,
|
|
fps=30,
|
|
ext="jpg",
|
|
down_scale=1,
|
|
start_sec=0,
|
|
end_sec=-1,
|
|
overwrite=False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
:param path
|
|
:param out_dir
|
|
:param fps
|
|
:param down_scale (optional int)
|
|
"""
|
|
path = str(path)
|
|
base_name = path.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
|
print(path, base_name)
|
|
print(os.path.isfile(path), os.path.isdir(base_name))
|
|
|
|
if os.path.isfile(path):
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
arg_str = f"-copyts -qscale:v 2 -vf fps={fps}"
|
|
if down_scale != 1:
|
|
arg_str = f"{arg_str},scale='iw/{down_scale}:ih/{down_scale}'"
|
|
if start_sec > 0:
|
|
arg_str = f"{arg_str} -ss {start_sec}"
|
|
if end_sec > start_sec:
|
|
arg_str = f"{arg_str} -to {end_sec}"
|
|
|
|
yn = "-y" if overwrite else "-n"
|
|
cmd = f"ffmpeg -i {path} {arg_str} {out_dir}/%06d.{ext} {yn}"
|
|
print(cmd)
|
|
elif os.path.isdir(base_name):
|
|
os.system(f"cp -r {base_name} {out_dir}")
|
|
return 0
|
|
else:
|
|
print(path, os.path.exists(path))
|
|
raise ValueError
|
|
|
|
return subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
|
|
|
# def video_to_frames(
|
|
# path,
|
|
# out_dir,
|
|
# fps=30,
|
|
# ext="jpg",
|
|
# down_scale=1,
|
|
# start_sec=0,
|
|
# end_sec=-1,
|
|
# overwrite=False,
|
|
# **kwargs,
|
|
# ):
|
|
# """
|
|
# :param path
|
|
# :param out_dir
|
|
# :param fps
|
|
# :param down_scale (optional int)
|
|
# """
|
|
# path = str(path)
|
|
# base_name = path.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
|
# print(path, base_name)
|
|
# print(os.path.isfile(path), os.path.isdir(base_name))
|
|
|
|
# if os.path.isfile(path):
|
|
# os.makedirs(out_dir, exist_ok=True)
|
|
|
|
# arg_str = f"-copyts -qscale:v 2 -vf fps={fps}"
|
|
# if down_scale != 1:
|
|
# arg_str = f"{arg_str},scale='iw/{down_scale}:ih/{down_scale}'"
|
|
# if start_sec > 0:
|
|
# arg_str = f"{arg_str} -ss {start_sec}"
|
|
# if end_sec > start_sec:
|
|
# arg_str = f"{arg_str} -to {end_sec}"
|
|
|
|
# yn = "-y" if overwrite else "-n"
|
|
# # Changed %06d to start from 0 instead of 1
|
|
# cmd = f"ffmpeg -i {path} {arg_str} -start_number 0 {out_dir}/%06d.{ext} {yn}"
|
|
# print(cmd)
|
|
# elif os.path.isdir(base_name):
|
|
# os.system(f"cp -r {base_name} {out_dir}")
|
|
# return 0
|
|
# else:
|
|
# print(path, os.path.exists(path))
|
|
# raise ValueError
|
|
|
|
# return subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--data_root", type=str, required=True, help="path to video")
|
|
parser.add_argument(
|
|
"--out_root", type=str, required=True, help="output dir for frames"
|
|
)
|
|
parser.add_argument(
|
|
"--seqs",
|
|
nargs="*",
|
|
default=None,
|
|
help="[optional] sequences to run, default runs all available",
|
|
)
|
|
parser.add_argument("--fps", type=int, default=30, help="fps to extract frames")
|
|
parser.add_argument(
|
|
"--ext", type=str, default="jpg", help="output filetype for frames"
|
|
)
|
|
parser.add_argument(
|
|
"--down_scale", type=int, default=1, help="scale to extract frames"
|
|
)
|
|
parser.add_argument(
|
|
"-ss", "--start_sec", type=float, default=0, help="seconds to start_sec"
|
|
)
|
|
parser.add_argument(
|
|
"-es", "--end_sec", type=float, default=-1, help="seconds to end_sec"
|
|
)
|
|
parser.add_argument(
|
|
"-y", "--overwrite", action="store_true", help="overwrite if already exist"
|
|
)
|
|
args = parser.parse_args()
|
|
seqs_all = os.listdir(args.data_root)
|
|
if args.seqs is None:
|
|
args.seqs = seqs_all
|
|
|
|
for seq in args.seqs:
|
|
path = os.path.join(args.data_root, seq)
|
|
print(f"EXTRACTING FRAMES FROM {path}")
|
|
assert os.path.isfile(path)
|
|
seq_name = os.path.splitext(os.path.basename(path.rstrip("/")))[0]
|
|
out_dir = os.path.join(args.out_root, seq_name)
|
|
video_to_frames(
|
|
path,
|
|
out_dir,
|
|
args.fps,
|
|
args.ext,
|
|
args.down_scale,
|
|
args.start_sec,
|
|
args.end_sec,
|
|
args.overwrite,
|
|
)
|