Sketchycoco：根据coco中图像获得对应的caption，并生成prompt.json文件

import os
import shutil
import json
import random


root = r"G:\SketchDiffusion\Sketchycoco\Dataset\Scene"
image_path = os.path.join(root, "source", "train")
filename_list = os.listdir(image_path)
prompts = []
for i, filename in enumerate(filename_list):
    img_name = filename.split('.')[0]
    # search txt matching with image
    caption_path = r"G:\SketchDiffusion\Sketchycoco\Dataset\caption"
    for caption_txt in os.listdir(caption_path):
        txt_name = caption_txt.split('.')[0]
        if txt_name == "coco_"+img_name:
            txt_path = os.path.join(caption_path, caption_txt)
            with open(txt_path, 'r') as f1:
                content = f1.readlines()
            row_index = random.randint(0, len(content) - 1)
            random_row = content[row_index].replace("\n", "")
            if random_row.endswith("."):
                random_row = random_row.replace(".", "")
            prompt = random_row + " in the style of sketch"
            prompts.append({"source": "source/train/"+filename,
                            "target": "target/train/"+img_name+'.png',
                            "prompt": prompt})
            break
    # if i > 4:
    #     break

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