blender-studio-tools/scripts/shotstats/shotstats.py

309 lines
9.5 KiB
Python
Executable File

#!/usr/bin/env python3
"""Generate visual render stats for a frame sequence.
Overlap a render time line chart to a video. Display a red playhead
corresponding to the current frame in the video.
This is some vintage code from 2017, written by fsiddi during the Agent 327
film project, to visually detect spikes in render time.
"""
import argparse
import csv
import datetime
import json
import os
import re
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile
parser = argparse.ArgumentParser(description='Generate shots stats.')
parser.add_argument('-i', '--in_path', help='Input path', required=True)
parser.add_argument('-o', '--out_path', help='Output directory')
parser.add_argument('-y', '--skip_confirmation', help='Skip confirmation', action='store_true')
parser.add_argument('-f', '--framerate', help='Framerate', default=24)
parser.add_argument('--image_format', help='Either PNG or EXR', default='PNG')
parser.add_argument('--memory_unit', default='G')
parser.add_argument('--render_time_unit', help='How display render time', default='m')
args = parser.parse_args()
def which(command):
"""Check if command is available and return its path."""
command_path = shutil.which(command)
if command_path is None:
print(f'{command} is required to run this script it, but it was not found.')
sys.exit()
return command_path
toolset = {
'ffmpeg': os.environ.get('FFMPEG_BIN', 'ffmpeg'),
'ffprobe': os.environ.get('FFPROBE_BIN', 'ffprobe'),
'exrheader': os.environ.get('EXRHEADER_BIN', 'exrheader'),
'gnuplot': os.environ.get('GNUPLOT_BIN', 'gnuplot'),
'identify': os.environ.get('IDENTIFY_BIN', 'identify'),
}
# Get render time and memory from frames
# If exr, use exrheader (we will also need jpeg previews later)
def parse_metadata(metadata_string: str):
if args.image_format == 'EXR':
# Metadata looks like this:
# Memory (type string): "0.00M"
re_result = re.search('"(.*)"', metadata_string)
return re_result.group(1)
elif args.image_format == 'PNG':
# Metadata looks like this:
# Memory: 0.00M
return metadata_string.split(':', 1)[1].strip()
def parse_memory(memory: str):
"""Get the amount of memory used.
We strip the last char, and assume it's M. Then we cast to float.
"""
memory = parse_metadata(memory)
memory_in_mb = float(memory[:-1])
if args.memory_unit == 'G':
m = memory_in_mb / 1024
else:
m = memory_in_mb
return m
def parse_render_time(time_metadata):
"""Get the render time in seconds."""
time_metadata = parse_metadata(time_metadata)
time_array = time_metadata.split(':')
if len(time_array) < 2: # Only seconds
time_in_seconds = float(time_metadata)
elif len(time_array) < 3: # Minutes and seconds
time_in_seconds = int(time_array[0]) * 60 + float(time_array[1])
elif len(time_array) < 4: # Hours, minutes and seconds
time_in_seconds = int(time_array[0]) * 3600 + int(time_array[1]) * 60 + float(time_array[2])
else:
time_in_seconds = float(0)
if args.render_time_unit == 'm':
t = time_in_seconds / 60
else:
t = time_in_seconds
return t
def parse_frame_number(frame_number_metadata):
"""Get the frame number."""
frame_number_metadata = parse_metadata(frame_number_metadata)
return int(frame_number_metadata)
def parse_exr_frames(frames_list):
"""Parse EXR frames using the 'exrheader' command."""
frames_stats = []
for frame in frames_list:
exrheader_command = [
toolset['exrheader'],
frame
]
p = subprocess.Popen(exrheader_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
out, err = p.communicate()
lines = iter(out.decode('utf-8').splitlines())
frame_stats = {
'name': frame.stem,
'frame_number': 0,
'memory_in_mb': 0,
'render_time_in_s': 0
}
for line in lines:
if line.startswith('Memory'):
frame_stats['memory_in_mb'] = parse_memory(line)
elif line.startswith('RenderTime'):
frame_stats['render_time_in_s'] = parse_render_time(line)
elif line.startswith('Frame'):
frame_stats['frame_number'] = parse_frame_number(line)
frames_stats.append(frame_stats)
return frames_stats
# If png use identify -verbose
def parse_png_frames(frames_list):
frames_stats = []
for frame in frames_list:
identify_command = [
toolset['identify'],
'-verbose',
frame
]
p = subprocess.Popen(identify_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
out, err = p.communicate()
lines = iter(out.decode('utf-8').splitlines())
frame_stats = {
'name': frame.stem,
'frame_number': 0,
'memory_in_mb': 0,
'render_time_in_s': 0
}
for line in lines:
line = line.strip()
if line.startswith('Memory'):
frame_stats['memory_in_mb'] = parse_memory(line)
elif line.startswith('RenderTime'):
frame_stats['render_time_in_s'] = parse_render_time(line)
# elif line.startswith('Frame'):
# frame_stats['frame_number'] = parse_frame_number(line)
frame_stats['frame_number'] = int(frame.stem)
print(f"Frame {frame_stats['frame_number']}: {frame_stats['memory_in_mb']} - {frame_stats['render_time_in_s']}")
frames_stats.append(frame_stats)
return frames_stats
# Get current directory
cwd = Path.cwd()
# Get absolute path of input dir (if relative it will be combined with cwd)
in_dir_absolute_path = cwd.joinpath(args.in_path)
# Look for files (png or exr)
frames = sorted(in_dir_absolute_path.glob(f'*.{args.image_format.lower()}'))
frames_stats_path = in_dir_absolute_path.parent / f'{in_dir_absolute_path.name}-frames_stats.csv'
def get_frame_stats():
if frames_stats_path.exists():
with open(frames_stats_path) as csvfile:
reader = csv.DictReader(csvfile, delimiter='\t')
return [row for row in reader]
if frames:
if args.image_format == 'EXR':
stats = parse_exr_frames(frames)
elif args.image_format == 'PNG':
stats = parse_png_frames(frames)
with open(frames_stats_path, 'w', newline='') as csvfile:
fieldnames = ['frame_number', 'name', 'memory_in_mb', 'render_time_in_s']
writer = csv.DictWriter(csvfile, fieldnames=fieldnames, delimiter='\t')
writer.writeheader()
for s in stats:
writer.writerow(s)
print(f'{frames_stats_path} is ready.')
return stats
else:
print(f'No {args.image_format} images found.')
sys.exit()
stats = get_frame_stats()
# Get frame resolution
# TODO(fsiddi) handle missing image
first_frame = frames[0]
# If we are working with exr, look for a .jpg file
if args.image_format == 'EXR':
first_frame = frames[0].with_suffix('.jpg')
identify_format = '%[fx:w]x%[fx:h]'
identify_command = [
'identify',
'-format',
identify_format,
first_frame
]
p = subprocess.Popen(identify_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
out, err = p.communicate()
result = out.decode('utf-8').splitlines()[0]
frame_width, frame_height = result.split('x')
# Make chart with memory usage and render time, using the size of frame
# tmp_dir = tempfile.TemporaryDirectory()
# tmp_dir_path = Path(tmp_dir.name)
gnuplot_chart_config_path = in_dir_absolute_path.parent / f'{in_dir_absolute_path.name}-gnuplot_chart'
chart_file_path = in_dir_absolute_path.parent / f'{in_dir_absolute_path.name}-chart.png'
template_vars = {
'tmp_chart_file': chart_file_path,
'frames_stats_file': frames_stats_path,
'frame_start_number': stats[0]['frame_number'],
'width': frame_width,
'height': frame_height,
}
with open('gnuplot_chart.tpl') as fp:
line = fp.readline()
with open(gnuplot_chart_config_path, 'w') as fc:
while line:
parsed_line = line.format(**template_vars)
fc.write(parsed_line)
line = fp.readline()
gnuplot_command = [
'gnuplot',
'-c',
gnuplot_chart_config_path,
]
subprocess.call(gnuplot_command)
#sys.exit()
# Combine the chart with images sequence and overlay the playhead
# For instance if we have a 260 frames clip
#
# 260 frames = 260 / <args.framerate> = 10.83333 seconds
# video_width / video_duration = 2048 / 10.83333 = 189.04621 pixels / second.
# Arbitrary offset defined by the chart
chart_margin_x_pixel = 45
chart_width = int(frame_width) - chart_margin_x_pixel - 17
pixel_per_second = chart_width / (len(frames) / args.framerate)
overlay_string = f"overlay, overlay=x='if(gte(t,0), -w+{chart_margin_x_pixel}+(t)*{pixel_per_second}, NAN)':y=0"
# Get the number of the first frame of the sequence
start_number = stats[0]['frame_number']
output_file_name = f'{in_dir_absolute_path.name}-{datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")}.mp4'
output_file_path = Path(args.out_path) / output_file_name
extension = 'png' if args.image_format == 'PNG' else 'jpg'
input_path = in_dir_absolute_path.joinpath(f'%6d.{extension}')
ffmpeg_command = [
'ffmpeg',
'-framerate',
f'{args.framerate}',
'-start_number',
f'{start_number}',
'-i',
f'{input_path}',
'-i',
f'{chart_file_path}',
'-i',
'playhead.png',
'-filter_complex',
f'{overlay_string}',
str(output_file_path)
]
subprocess.call(ffmpeg_command)