python pillow从一批图中随机选一部分粘贴到另一批图中的特定位置

随机选一部分头，对于每个头，从所有人中随机选一个粘贴到他头上

from PIL import Image
import os
import random
import json
import shutil
import numpy as np

# 获取头的坐标位置
def getHead(label_dir,label_file):
    anno = json.load(open(label_dir+label_file, "r", encoding="utf-8"))
    # 整张图的宽高
    w0, h0 = anno['imageWidth'], anno['imageHeight']
    shapes = anno['shapes']
    pts = shapes[0]['points']
    x1 = max(0,int(pts[0][0]))
    y1 = max(0,int(pts[0][1]))
    x2 = min(w0,int(pts[1][0]))
    y2 = min(h0,int(pts[1][1]))
    # 头的宽高
    w  = abs(x2 - x1)
    h  = abs(y2 - y1)
    return w, h,w0,h0,x1,y1,x2,y2


# 检查是不是只有一个人
def checkPeopleNum(txt_path):
    with open(txt_path) as f:
            mask = np.array(list(map(lambda x:np.array(x.strip().split()), f.readlines())))
    if len(mask) > 1:
        return False
    return True


def handle_img(imgBig,imgSmall,imgSave,nums):
    os.makedirs(imgSave,exist_ok=True)

    # 所有只有头的图片,随机打乱顺序之后取一部分
    imgSList = os.listdir(imgSmall)
    random.shuffle(imgSList)
    print('len(imgSList):',len(imgSList))
    imgSList = imgSList[:nums]

    # 所有人的图片，不用打乱顺序也不用取一部分，因为是对于每张只有头的图片，都从人的列表中选一个来paste
    jsonBList = [i for i in os.listdir(imgBig) if i.endswith('.json')]
    # random.shuffle(jsonBList)
    print('len(jsonBList):',len(jsonBList))
    # jsonBList = jsonBList[:10]

    # 相当于随机选一部分头，对于每个头，随机粘贴到某个人的头上
    for i,imgSP in enumerate(imgSList):
        jsonBP = random.choice(jsonBList) # 从人的列表中随机选一个
        # 只粘贴到只有一个人的图上
        while not checkPeopleNum(imgBig + jsonBP.split('.json')[0] + '.txt'):
            jsonBP = random.choice(jsonBList)

        print(jsonBP,imgSP,'path')
        w,h,w0,h0,x1,y1,x2,y2 = getHead(imgBig,jsonBP)
        imgS = Image.open(imgSmall+imgSP)
        imgS = imgS.resize((w,h))

        imgBP = jsonBP.split('.json')[0]+'.jpg'
        imgB = Image.open(imgBig+imgBP)

        imgB.paste(imgS, (min(x1,x2),min(y1,y2)))
        imgB = imgB.convert('RGB')
        imgB.save(imgSave+imgBP[:-4]+'_'+str(i)+'.jpg')

        # 因为都是用的头盔粘贴到人头上,所以label一定是helmet
        with open(imgBig+imgBP[:-4]+'.txt') as f:
            mask = np.array(list(map(lambda x:np.array(x.strip().split()), f.readlines())))
        cls_id = int(float(mask[0][0]))
        if cls_id != 0:
            # print(imgSave+imgBP[:-4]+'_'+str(i)+'.txt')
            mask[0][0] = 0.0
            # print(mask[0])
            with open(imgSave+imgBP[:-4]+'_'+str(i)+'.txt','w') as f:
                f.write(" ".join(mask[0]))
        else:
            shutil.copy(imgBig+imgBP[:-4]+'.txt', imgSave+imgBP[:-4]+'_'+str(i)+'.txt')


if __name__ == '__main__':
    imgBig = '../data/peoples/'
    imgSmall = '../data/helmet_classify20230324/val/helmet/'
    imgSave = '../data/peopleAug/'
    nums = 100
    handle_img(imgBig,imgSmall,imgSave,nums)