File size: 1,761 Bytes
f927ea3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | from pathlib import Path
import shutil,json,hashlib
from PIL import Image
root=Path('/data/nyako-training');src=root/'source_dataset';dst=root/'artifacts/dataset';dst.mkdir(exist_ok=True)
prefix=b'@nyako, '; rows=[]
for p in sorted(src.rglob('*.jpg')):
c=p.with_suffix('.txt'); assert c.exists(), 'Missing caption'
rel=p.relative_to(src); target=dst/rel;target.parent.mkdir(parents=True,exist_ok=True);shutil.copy2(p,target)
# Opaque byte operation only: caption text is never printed or interpreted.
original=c.read_bytes(); amended=prefix+original
target.with_suffix('.txt').write_bytes(amended)
assert target.with_suffix('.txt').read_bytes()[len(prefix):]==original
with Image.open(p) as im: size=im.size
rows.append({'image':str(rel),'size':size,'original_caption_sha256':hashlib.sha256(original).hexdigest(),'prefixed_caption_sha256':hashlib.sha256(amended).hexdigest()})
assert rows
(root/'artifacts/config/dataset_manifest.json').write_text(json.dumps({'repo':'ACCC1380/Nyako','revision':'f1048ee38473bc253ebdc4b4eba3b0c2dcebacb3','prefix':'@nyako, ','count':len(rows),'files':rows},indent=2))
print('Prepared image/caption pairs:',len(rows)); print('Image dimensions range:',[min(r['size'][i] for r in rows) for i in (0,1)],[max(r['size'][i] for r in rows) for i in (0,1)])
folders=sorted(set(str((dst/r['image']).parent) for r in rows))
text='[general]\ncaption_extension = ".txt"\nshuffle_caption = false\n\n[[datasets]]\nresolution = 1024\nbatch_size = 2\nenable_bucket = true\nmin_bucket_reso = 256\nmax_bucket_reso = 2048\nbucket_reso_steps = 64\nbucket_no_upscale = false\n'
for folder in folders:text+='\n[[datasets.subsets]]\nimage_dir = '+json.dumps(folder)+'\nnum_repeats = 1\n'
(root/'artifacts/config/dataset.toml').write_text(text)
|