File size: 2,739 Bytes
00c7586 | 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 hashlib, io
from PIL import Image
PACKED_IMAGE_LIMIT_BYTES = 21_866_496
SCALES = [1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2]
def sha(b):
return hashlib.sha256(b).hexdigest()
def png(im):
b = io.BytesIO()
im.save(b, format='PNG', compress_level=1, optimize=False)
return b.getvalue()
def normalized(plan, frames, expected=None):
factors = [expected['common_geometry']['scale']] if expected else SCALES
for factor in factors:
blocks = []
count = 0
for block in plan['blocks']:
assets = []
for start in range(0, len(block['indices']), 2):
indices = block['indices'][start:start + 2]
originals = [frames[i] for i in indices]
w, h = originals[0].size
nw = round(w * factor)
nh = round(h * factor)
packed = Image.new('RGB', (nw * len(indices), nh))
rgb = []
for pos, image in enumerate(originals):
if factor != 1:
image = image.resize((nw, nh), Image.Resampling.LANCZOS)
packed.paste(image, (pos * nw, 0))
rgb.append(sha(image.tobytes()))
blob = png(packed)
count += len(blob) + (-len(blob) % 3)
assets.append(dict(indices=indices, size=list(packed.size), sha256=sha(blob), rgb_sha256=rgb, bytes=len(blob), _blob=blob))
blocks.append(dict(role=block['role'], images=assets))
if count <= PACKED_IMAGE_LIMIT_BYTES:
break
else:
raise RuntimeError('No common caption input geometry fits')
if expected:
for old, new in zip(expected['frame_blocks'], blocks):
assert old['role'] == new['role'] and len(old['images']) == len(new['images'])
for a, b in zip(old['images'], new['images']):
assert a['indices'] == b['indices'] and a['rgb_sha256'] == b['rgb_sha256'] and (a['size'] == b['size'])
return dict(plan, frame_blocks=blocks, common_geometry=dict(scale=factor, source_pixel_unchanged=factor == 1, aligned_image_bytes=count, image_size_limit_bytes=PACKED_IMAGE_LIMIT_BYTES, sampling_frames_removed=0))
def source_frames(normalized):
out = {}
for block in normalized['frame_blocks']:
for a in block['images']:
image = Image.open(io.BytesIO(a['_blob'])).convert('RGB')
n = len(a['indices'])
w = image.width // n
for pos, i in enumerate(a['indices']):
frame = image.crop((pos * w, 0, (pos + 1) * w, image.height))
assert sha(frame.tobytes()) == a['rgb_sha256'][pos]
out[i] = frame
return out
|