283 lines
10 KiB
Python
283 lines
10 KiB
Python
#(c) Alex Spirin 2023
|
|
|
|
import hashlib, os, sys, glob, subprocess, pathlib, platform
|
|
import zipfile
|
|
from shutil import copy
|
|
import requests
|
|
|
|
def generate_file_hash(input_file):
|
|
"""Generates has for video vile based on name, size, creation time
|
|
"""
|
|
# Get file name and metadata
|
|
file_name = os.path.basename(input_file)
|
|
file_size = os.path.getsize(input_file)
|
|
creation_time = os.path.getctime(input_file)
|
|
|
|
# Generate hash
|
|
hasher = hashlib.sha256()
|
|
hasher.update(file_name.encode('utf-8'))
|
|
hasher.update(str(file_size).encode('utf-8'))
|
|
hasher.update(str(creation_time).encode('utf-8'))
|
|
file_hash = hasher.hexdigest()
|
|
|
|
return file_hash
|
|
|
|
def createPath(filepath):
|
|
os.makedirs(filepath, exist_ok=True)
|
|
|
|
def extractFrames(video_path, output_path, nth_frame, start_frame, end_frame):
|
|
createPath(output_path)
|
|
print(f"Exporting Video Frames (1 every {nth_frame})...")
|
|
try:
|
|
for f in [o.replace('\\','/') for o in glob.glob(output_path+'/*.jpg')]:
|
|
pathlib.Path(f).unlink()
|
|
except:
|
|
print('error deleting frame ', f)
|
|
vf = f'select=between(n\\,{start_frame}\\,{end_frame}) , select=not(mod(n\\,{nth_frame}))'
|
|
if os.path.exists(video_path):
|
|
try:
|
|
subprocess.run(['ffmpeg', '-i', f'{video_path}', '-vf', f'{vf}', '-vsync', 'vfr', '-q:v', '2', '-loglevel', 'error', '-stats', f'{output_path}/%06d.jpg'], stdout=subprocess.PIPE).stdout.decode('utf-8')
|
|
except:
|
|
subprocess.run(['ffmpeg.exe', '-i', f'{video_path}', '-vf', f'{vf}', '-vsync', 'vfr', '-q:v', '2', '-loglevel', 'error', '-stats', f'{output_path}/%06d.jpg'], stdout=subprocess.PIPE).stdout.decode('utf-8')
|
|
|
|
else:
|
|
sys.exit(f'\nERROR!\n\nVideo not found: {video_path}.\nPlease check your video path.\n')
|
|
|
|
|
|
class FrameDataset():
|
|
def __init__(self, source_path, outdir_prefix='', videoframes_root='', update_on_getitem=False, start_frame=0, end_frame=-1, nth_frame=1, overwrite=False):
|
|
if outdir_prefix == '':
|
|
outdir_prefix = f'{start_frame}_{end_frame}_{nth_frame}'
|
|
if end_frame == -1: end_frame = 999999999
|
|
self.frame_paths = None
|
|
image_extenstions = ['jpeg', 'jpg', 'png', 'tiff', 'bmp', 'webp']
|
|
self.update_on_getitem = update_on_getitem
|
|
self.source_path = source_path
|
|
if not os.path.exists(source_path):
|
|
if len(glob.glob(source_path))>0:
|
|
"""if non-empty glob-pattern"""
|
|
self.frame_paths = sorted(glob.glob(source_path))[start_frame:end_frame:nth_frame]
|
|
else:
|
|
raise FileNotFoundError(f'Frame source for {outdir_prefix} not found at {source_path}\nPlease specify an existing source path.')
|
|
if os.path.exists(source_path):
|
|
if os.path.isfile(source_path):
|
|
if os.path.splitext(source_path)[1][1:].lower() in image_extenstions:
|
|
"""if 1 image"""
|
|
self.frame_paths = [source_path]
|
|
else:
|
|
"""if 1 video"""
|
|
hash = generate_file_hash(source_path)[:10]
|
|
out_path = os.path.join(videoframes_root, outdir_prefix+'_'+hash)
|
|
files = glob.glob(os.path.join(out_path, '*.*'))
|
|
if len(files)>0 and not overwrite:
|
|
self.frame_paths = files
|
|
print(f'Found {len(self.frame_paths)} frames in {out_path}. Skipping extraction. Check overwrite option to overwrite.')
|
|
return
|
|
else:
|
|
print(f'Extracting frames from {source_path} to {out_path}')
|
|
extractFrames(source_path, out_path,
|
|
nth_frame=nth_frame, start_frame=start_frame, end_frame=end_frame)
|
|
self.frame_paths = glob.glob(os.path.join(out_path, '*.*')) #dont apply start-end here as already applied during video extraction
|
|
self.source_path = out_path
|
|
if len(self.frame_paths)<1:
|
|
raise FileNotFoundError(f'Couldn`t extract frames from {source_path}\nPlease specify an existing source path.')
|
|
elif os.path.isdir(source_path):
|
|
"""if folder with images"""
|
|
self.frame_paths = glob.glob(os.path.join(source_path, '*.*'))[start_frame:end_frame:nth_frame]
|
|
if len(self.frame_paths)<1:
|
|
raise FileNotFoundError(f'Found 0 frames in {source_path}\nPlease specify an existing source path.')
|
|
extensions = []
|
|
if self.frame_paths is not None:
|
|
for f in self.frame_paths:
|
|
ext = os.path.splitext(f)[1][1:]
|
|
if ext not in image_extenstions:
|
|
raise Exception(f'Found non-image file extension: {ext} in {source_path}. Please provide a folder with image files of the same extension, or specify a glob pattern.')
|
|
if not ext in extensions:
|
|
extensions+=[ext]
|
|
if len(extensions)>1:
|
|
raise Exception(f'Found multiple file extensions: {extensions} in {source_path}. Please provide a folder with image files of the same extension, or specify a glob pattern.')
|
|
|
|
self.frame_paths = sorted(self.frame_paths)
|
|
|
|
else: raise FileNotFoundError(f'Frame source for {outdir_prefix} not found at {source_path}\nPlease specify an existing source path.')
|
|
print(f'Found {len(self.frame_paths)} frames at {source_path}')
|
|
|
|
def __getitem__(self, idx):
|
|
if self.update_on_getitem:
|
|
self.frame_paths = glob.glob(os.path.join(self.source_path, '*.*'))
|
|
|
|
idx = min(idx, len(self.frame_paths)-1)
|
|
return self.frame_paths[idx]
|
|
|
|
def __len__(self):
|
|
return len(self.frame_paths)
|
|
|
|
|
|
class StylizedFrameDataset():
|
|
def __init__(self, source_path):
|
|
self.frame_paths = None
|
|
self.source_path = source_path
|
|
if not os.path.exists(source_path):
|
|
raise FileNotFoundError(f'Frame source not found at {source_path}\nPlease specify an existing source path.')
|
|
if os.path.exists(source_path):
|
|
if os.path.isfile(source_path):
|
|
raise NotADirectoryError(f'{source_path} is a file. Please specify path to a folder.')
|
|
elif os.path.isdir(source_path):
|
|
self.frame_paths = glob.glob(os.path.join(source_path, '*.*'))
|
|
|
|
def __getitem__(self, idx):
|
|
self.frame_paths = glob.glob(os.path.join(self.source_path, '*.*'))
|
|
if len(self.frame_paths)==0:
|
|
return ''
|
|
|
|
idx = min(idx, len(self.frame_paths)-1)
|
|
return self.frame_paths[idx]
|
|
|
|
def __len__(self):
|
|
return len(self.frame_paths)
|
|
|
|
|
|
def get_size(size, max_size, divisible_by=8):
|
|
divisible_by = int(divisible_by)
|
|
x,y = size
|
|
max_dim = max(size)
|
|
# if max_dim>max_size:
|
|
ratio = max_size/max_dim
|
|
new_size = ((int(x*ratio)//divisible_by*divisible_by),(int(y*ratio)//divisible_by*divisible_by))
|
|
return new_size
|
|
# return size
|
|
|
|
def find_ffmpeg(start_dir):
|
|
files = glob.glob(f"{start_dir}/**/*.*", recursive=True)
|
|
for f in files:
|
|
if platform.system() == 'Linux':
|
|
if f.endswith('ffmpeg'): return f
|
|
elif f.endswith('ffmpeg.exe'): return f
|
|
return None
|
|
|
|
def download_ffmpeg(start_dir):
|
|
ffmpeg_url = 'https://github.com/GyanD/codexffmpeg/releases/download/6.0/ffmpeg-6.0-full_build.zip'
|
|
print('ffmpeg.exe not found, downloading...')
|
|
r = requests.get(ffmpeg_url, allow_redirects=True, timeout=60)
|
|
print('downloaded, extracting')
|
|
open('ffmpeg-6.0-full_build.zip', 'wb').write(r.content)
|
|
|
|
with zipfile.ZipFile('ffmpeg-6.0-full_build.zip', 'r') as zip_ref:
|
|
zip_ref.extractall(f'{start_dir}/')
|
|
|
|
copy(f'{start_dir}/ffmpeg-6.0-full_build/bin/ffmpeg.exe', f'{start_dir}/')
|
|
return f'{start_dir}/ffmpeg.exe'
|
|
|
|
|
|
def get_ffmpeg():
|
|
start_dir = os.getcwd()
|
|
ffmpeg_path = find_ffmpeg(start_dir)
|
|
if ffmpeg_path is None or not os.path.exists(ffmpeg_path):
|
|
ffmpeg_path = download_ffmpeg(start_dir)
|
|
return ffmpeg_path
|
|
|
|
def save_video(indir, video_out, batch_name='', start_frame=1, last_frame=-1, fps=30, output_format='h264_mp4', use_deflicker=False):
|
|
ffmpeg_path = get_ffmpeg()
|
|
print('Found ffmpeg at: ', ffmpeg_path)
|
|
os.makedirs(video_out, exist_ok=True)
|
|
indir = indir.replace('\\','/')
|
|
image_path = f"{indir}/{batch_name}_%06d.png"
|
|
|
|
postfix = ''
|
|
if use_deflicker:
|
|
postfix+='_dfl'
|
|
|
|
indir_stem = indir.replace('\\','/').strip('/').split('/')[-1]
|
|
out_filepath = os.path.join(video_out, f"{indir_stem}_{batch_name}_{postfix}.{output_format.split('_')[-1]}")
|
|
if last_frame == -1:
|
|
last_frame = len(glob.glob(f"{indir}/*.png"))
|
|
|
|
cmd = [ffmpeg_path,
|
|
'-y',
|
|
'-vcodec',
|
|
'png',
|
|
'-r',
|
|
str(fps),
|
|
'-start_number',
|
|
str(start_frame),
|
|
'-i',
|
|
image_path,
|
|
'-frames:v',
|
|
str(last_frame+1),
|
|
'-c:v']
|
|
|
|
if output_format == 'h264_mp4':
|
|
cmd+=['libx264',
|
|
'-pix_fmt',
|
|
'yuv420p']
|
|
elif output_format == 'qtrle_mov':
|
|
cmd+=['qtrle',
|
|
'-vf',
|
|
f'fps={fps}']
|
|
elif output_format == 'prores_mov':
|
|
cmd+=['prores_aw',
|
|
'-profile:v',
|
|
'2',
|
|
'-pix_fmt',
|
|
'yuv422p10',
|
|
'-vf',
|
|
f'fps={fps}']
|
|
|
|
if use_deflicker:
|
|
cmd+=['-vf','deflicker=mode=pm:size=10']
|
|
cmd+=[out_filepath]
|
|
|
|
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
stdout, stderr = process.communicate()
|
|
if process.returncode != 0:
|
|
print(stderr)
|
|
raise RuntimeError(stderr)
|
|
else:
|
|
print(f"The video is ready and saved to {out_filepath}")
|
|
return 0
|
|
|
|
def get_sched_from_json(frame_num, sched_json, blend=False):
|
|
frame_num = int(frame_num)
|
|
frame_num = max(frame_num, 0)
|
|
sched_int = {}
|
|
for key in sched_json.keys():
|
|
sched_int[int(key)] = sched_json[key]
|
|
sched_json = sched_int
|
|
keys = sorted(list(sched_json.keys()))
|
|
if frame_num<0:
|
|
frame_num = max(keys)
|
|
try:
|
|
frame_num = min(frame_num,max(keys)) #clamp frame num to 0:max(keys) range
|
|
except:
|
|
pass
|
|
|
|
if frame_num in keys:
|
|
return sched_json[frame_num]
|
|
if frame_num not in keys:
|
|
for i in range(len(keys)-1):
|
|
k1 = keys[i]
|
|
k2 = keys[i+1]
|
|
if frame_num > k1 and frame_num < k2:
|
|
if not blend:
|
|
print('frame between keys, no blend')
|
|
return sched_json[k1]
|
|
if blend:
|
|
total_dist = k2-k1
|
|
dist_from_k1 = frame_num - k1
|
|
return sched_json[k1]*(1 - dist_from_k1/total_dist) + sched_json[k2]*(dist_from_k1/total_dist)
|
|
return 0
|
|
|
|
def get_scheduled_arg(frame_num, schedule, blend_json_schedules=True):
|
|
if ':' in schedule:
|
|
schedule = eval("{"+schedule+"}")
|
|
else:
|
|
schedule = eval(str(schedule))
|
|
print(schedule)
|
|
if isinstance(schedule, int):
|
|
return schedule
|
|
if isinstance(schedule, float):
|
|
return schedule
|
|
if isinstance(schedule, list):
|
|
return schedule[frame_num] if frame_num<len(schedule) else schedule[-1]
|
|
if isinstance(schedule, dict):
|
|
return get_sched_from_json(frame_num, schedule, blend=blend_json_schedules) |