import argparse
import csv
import os
import subprocess
import tempfile
import numpy
from collections import defaultdict
def run_command(command):
"""Run a shell command and return its output."""
result = subprocess.run(command, capture_output=True, text=True, shell=True)
if result.returncode != 0:
raise subprocess.CalledProcessError(result.returncode, command, result.stderr)
return result.stdout.strip()
def encode_svt_av1(input_file, preset: int, crf: int):
"""Encode input file with SVT-AV1."""
with tempfile.NamedTemporaryFile(suffix='.avif', delete=False) as temp_avif:
avif_path = temp_avif.name
full_command = f"avifenc -s {preset} -c svt -y 420 -d 10 -a \"crf={crf}\" -a tune=4 {input_file} -o {avif_path}"
run_command(full_command)
return avif_path
def encode_aomenc(input_file, preset: int, crf: int):
"""Encode input file with aomenc."""
with tempfile.NamedTemporaryFile(suffix='.avif', delete=False) as temp_avif:
avif_path = temp_avif.name
full_command = f"avifenc -j 8 -d 10 -y 444 -s {preset} --min 0 --max 63 --minalpha 0 --maxalpha 63 -a end-usage=q -a cq-level={crf} -a tune=ssim -a quant-b-adapt=1 -a deltaq-mode=2 -a sb-size=64 -a sharpness=1 {input_file} -o {avif_path}"
run_command(full_command)
return avif_path
def encode_jpegli(input_file, quality: int):
"""Encode input file with jpegli."""
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_jpg:
jpg_path = temp_jpg.name
full_command = f"cjpegli {input_file} {jpg_path} -q {quality} -p 2"
run_command(full_command)
return jpg_path
def encode_jxl(input_file, effort: int, quality: int):
with tempfile.NamedTemporaryFile(suffix='.jxl', delete=False) as temp_jxl:
jxl_path = temp_jxl.name
full_command = f"cjxl {input_file} {jxl_path} -q {quality} -e {effort}"
run_command(full_command)
return jxl_path
def encode_svt_av1_ffmpeg(input_file, preset, crf):
"""Encode input file with SVT-AV1 using ffmpeg."""
with tempfile.NamedTemporaryFile(suffix='.avif', delete=False) as temp_avif:
avif_path = temp_avif.name
full_command = f"ffmpeg -y -hide_banner -loglevel error -i \"{input_file}\" -pix_fmt yuv420p10le -vf scale=in_range=full:out_range=full -strict -2 -f yuv4mpegpipe - | ./SvtAv1EncApp -i - -b - --preset {preset} --crf {crf} --tune 4 --progress 3 --color-range 1 --use-fixed-qindex-offsets 1 | ffmpeg -y -hide_banner -loglevel error -i - -c copy \"{avif_path}\""
run_command(full_command)
return avif_path
def encode_x264(input_file, preset: str, crf: int):
"""Encode input file with x264."""
with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_mp4:
mp4_path = temp_mp4.name
full_command = f"ffmpeg -y -hide_banner -loglevel error -i \"{input_file}\" -pix_fmt yuv444p -frames:v 1 -c:v libx264 -preset {preset} -crf {crf} -tune ssim \"{mp4_path}\""
run_command(full_command)
return mp4_path
def convert_to_png(input_file):
"""Convert input file to PNG using avifdec."""
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_png:
png_path = temp_png.name
run_command(f"avifdec \"{input_file}\" \"{png_path}\" --png-compress 0")
return png_path
def convert_to_png_magick(input_file):
"""Convert input file to PNG using magick."""
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_png:
png_path = temp_png.name
run_command(f"magick -quality 0 \"{input_file}\" \"{png_path}\"")
return png_path
def calculate_ssimulacra2(original, encoded):
"""Calculate SSIMULACRA2 score."""
return float(run_command(f"ssimulacra2 \"{original}\" \"{encoded}\""))
def process_image(input_file, output_csv):
"""Process a single image and generate CSV with SSIMULACRA2 scores."""
results = []
# SVT-AV1 4:2:0
for i in range(0, 61, 5):
print(f"Encoding {os.path.basename(input_file)} with SVT-AV1-PSY CRF {i}...")
avif_path = encode_svt_av1(input_file, preset=4, crf=i)
# avif_path = encode_svt_av1_ffmpeg(ref_mkv_path, preset=4, crf=i)
png_path = convert_to_png(avif_path)
size = os.path.getsize(avif_path)
ssimulacra2_score = calculate_ssimulacra2(input_file, png_path)
results.append(['chroma_qm1', i, size, ssimulacra2_score])
os.unlink(avif_path)
os.unlink(png_path)
# aomenc 4:4:4
for i in range(5, 61, 5):
print(f"Encoding {os.path.basename(input_file)} with aomenc 4:4:4 CRF {i}...")
aom_path = encode_aomenc(input_file, preset=4, crf=i)
png_path = convert_to_png(aom_path)
size = os.path.getsize(aom_path)
ssimulacra2_score = calculate_ssimulacra2(input_file, png_path)
results.append(['AOM_444', i, size, ssimulacra2_score])
os.unlink(aom_path)
os.unlink(png_path)
# libx264 4:4:4
for i in range(0, 51, 5):
print(f"Encoding {os.path.basename(input_file)} with x264 4:4:4 CRF {i}...")
x264_path = encode_x264(input_file, preset="veryslow", crf=i)
png_path = convert_to_png_magick(x264_path)
size = os.path.getsize(x264_path)
ssimulacra2_score = calculate_ssimulacra2(input_file, png_path)
results.append(['X264_444', i, size, ssimulacra2_score])
os.unlink(x264_path)
os.unlink(png_path)
# cjpegli
for i in range(5, 100, 5):
print(f"Encoding {os.path.basename(input_file)} with jpegli quality {i}...")
jpg_path = encode_jpegli(input_file, quality=i)
size = os.path.getsize(jpg_path)
ssimulacra2_score = calculate_ssimulacra2(input_file, jpg_path)
results.append(['JPEGLI', i, size, ssimulacra2_score])
os.unlink(jpg_path)
# jpeg-xl
for i in range(5, 99, 5):
print(f"Encoding {os.path.basename(input_file)} with cjxl quality {i}...")
jxl_path = encode_jxl(input_file, effort=10, quality=i)
size = os.path.getsize(jxl_path)
ssimulacra2_score = calculate_ssimulacra2(input_file, jxl_path)
results.append(['JXL', i, size, ssimulacra2_score])
os.unlink(jxl_path)
with open(output_csv, 'w', newline='') as csvfile:
csvwriter = csv.writer(csvfile)
csvwriter.writerow(['codec', 'quality', 'size', 'ssimulacra2'])
csvwriter.writerows(results)
return results
def process_directory(input_dir, output_dir):
"""Process all images in the input directory and generate CSV files."""
os.makedirs(output_dir, exist_ok=True)
all_results = []
for filename in os.listdir(input_dir):
if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp')):
input_path = os.path.join(input_dir, filename)
output_csv = os.path.join(output_dir, f"{os.path.splitext(filename)[0]}_results.csv")
results = process_image(input_path, output_csv)
all_results.extend(results)
print(f"Results for {filename} saved to {output_csv}")
return all_results
def calculate_stats(all_results):
"""Calculate statistics of all results."""
processed_results = defaultdict(list)
for result in all_results:
codec, quality, size, score = result
key = (codec, quality)
processed_results[key].append((size, score))
return [
[codec, quality, int(numpy.average(list(zip(*value))[0])), numpy.average(list(zip(*value))[1]), numpy.percentile(list(zip(*value))[1], 10)]
for (codec, quality), value in processed_results.items()
]
def save_stats(average_results, output_file):
"""Save stats to a CSV file."""
with open(output_file, 'w', newline='') as csvfile:
csvwriter = csv.writer(csvfile)
csvwriter.writerow(['codec', 'quality', 'avg_size', 'avg_ssimulacra2', '10pct_ssimulacra2'])
csvwriter.writerows(average_results)
def main():
parser = argparse.ArgumentParser(description="Generate SSIMULACRA2 scores for different codecs")
parser.add_argument("input_dir", help="Input directory containing image files")
parser.add_argument("output_dir", help="Output directory for CSV files")
args = parser.parse_args()
all_results = process_directory(args.input_dir, args.output_dir)
stat_results = calculate_stats(all_results)
overall_average_file = os.path.join(args.output_dir, "0stats.csv")
save_stats(stat_results, overall_average_file)
print(f"Stats saved to {overall_average_file}")
if __name__ == "__main__":
main()