File size: 2,896 Bytes
3088c65 | 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 | import subprocess
import json
import os
import torch
import shutil
import server
import folder_paths
web = server.web
@server.PromptServer.instance.routes.post("/VHS_test")
async def test(request):
try:
req_data = await request.json()
output = req_data['output']['gifs'][0]
filename = output['filename']
typ = output['type']
base_args = ["ffprobe", "-v", "error", '-count_packets', "-show_entries", "stream", "-of", "json"]
video = folder_paths.get_annotated_filepath(f'{filename} [{typ}]')
vprobe = json.loads(subprocess.run(base_args + ['-select_streams', 'v:0', video],
capture_output=True, check=True).stdout)['streams'][0]
aprobe = json.loads(subprocess.run(base_args + ['-select_streams', 'a:0', video],
capture_output=True, check=True).stdout)['streams']
probe = {'video': vprobe}
if len(aprobe) > 0:
probe['audio'] = aprobe[0]
errors = []
compare = None
for test in req_data['tests']:
if test['type'] == 'compare':
compare = test
continue
key = test['key']
expected = test['value']
actual = probe[test['type']][key]
if expected != actual:
#Consider always dumping type?
errors.append(f'{key}: {expected} != {actual}')
if len(errors) == 0 and compare is not None:
if not os.path.exists(compare['filename']):
os.makedirs(os.path.split(compare['filename'])[0], exist_ok=True)
shutil.copy(video, compare['filename'])
print("Missing comparison file has been initialized from output:", os.path.abspath(compare['filename']))
else:
#NOTE: This does not include the full memory optimizations of VHS
#Tests should be small
#TODO: Figure out way to do opacity comparison. May need to do blending in python
#(easy, but slower and more memory intensive)
diff = subprocess.run(['ffmpeg', '-v', 'error', '-i', video, '-i', compare['filename'], '-filter_complex', 'blend=all_mode=grainextract', '-pix_fmt', 'rgb24', '-f', 'rawvideo', '-'], stdout=subprocess.PIPE, check=True).stdout
diff = torch.frombuffer(diff, dtype=torch.uint8).to(dtype=torch.float32).div_(255)
#diff = diff.reshape((-1,4))
d = (diff-0.5).abs().sum()/diff.size(0)
if d > compare['tolerance']:
errors.append(f'Similarity is outside specified tolerance: {d}')
else:
print('d:', d)
return web.json_response(errors)
except Exception as e:
return web.json_response(str(e))
|