SketchyCOCO数据集中Object创建prompt

import os
import shutil
import json
from operator import itemgetter

style = {2:'bicycle', 3:'car', 4:'motorcycle', 5:'airplane', 10:'traffic light',
         11:'fire hydrant', 17:'cat', 18:'dog', 19:'horse', 20:'sheep', 21:'cow',
         22:'elephant', 24:'zebra', 25:'giraffe'}

split = "val"
root = r"G:\SketchDiffusion\Sketchycoco\Dataset\Object"
image_path = os.path.join(root, "source", split)
imagedir_list = os.listdir(image_path)

prompts = []
for i, imagedir in enumerate(imagedir_list):
    stylename = style[int(imagedir)]
    prompt = 'a picture of ' + stylename + ' in the style of sketch'
    imagedir = os.path.join(image_path, imagedir)
    for filename in os.listdir(imagedir):
        img_name = filename.split('.')[0]
        prompts.append({"source": "source/"+split+"/" + filename,
                        "target": "target/"+split+"/" + img_name + '.png',
                        "prompt": prompt})
        srcfile = os.path.join(imagedir, filename)
        dstfile = os.path.join(image_path, filename)
        shutil.move(srcfile, dstfile)

prompts = sorted(prompts, key=itemgetter('source'))


json_path = os.path.join(root, split+"_prompt.json")
with open(json_path, 'w') as f2:
    for line in prompts:
        json.dump(line, f2)
        f2.write('\n')
print('over！')