| """Decode untrusted entropy payloads outside the web-server process.""" |
| import csv |
| import json |
| import resource |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
| import backend |
| from bitstream import parse_bitstream |
|
|
|
|
| def main(): |
| resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) |
| directory = Path(sys.argv[1]) |
| data = (directory / 'reconstruction.afp').read_bytes() |
| strings, info = parse_bitstream(data) |
| model = backend.get_model() |
| from src.utils import img_utils |
| with torch.no_grad(): |
| if backend.DEVICE.startswith('cuda'): |
| torch.cuda.synchronize() |
| start = time.perf_counter() |
| decoded, _, _ = backend.decode_strings(model, strings) |
| if backend.DEVICE.startswith('cuda'): |
| torch.cuda.synchronize() |
| elapsed = (time.perf_counter()-start)*1000 |
| if tuple(decoded.shape[-2:]) != (info['height'], info['width']) or not torch.isfinite(decoded).all(): |
| raise ValueError('Invalid reconstruction.') |
| img_utils.imwrite(str(directory/'reconstruction.png'), decoded) |
| result = dict(dataset='user_bitstream', image_name='reconstruction.png', mode='decode', |
| **info, nominal_bpp=backend.TARGETS[info['quality_index']], |
| bpp=len(data)*8/(info['width']*info['height']), bitstream_bytes=len(data), |
| decode_ms=elapsed, device=backend.DEVICE, |
| metric_protocol='No reference image supplied; PSNR and SSIM are not computed.', |
| timing_note='Single-request decoder timing excludes model loading.') |
| (directory/'metrics.json').write_text(json.dumps(result,indent=2)) |
| with (directory/'metrics.csv').open('w',newline='') as f: |
| w=csv.DictWriter(f,fieldnames=list(result));w.writeheader();w.writerow(result) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
