python-合成图

import os
import numpy as np
import cv2

def find_boundRect(mask):
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # 检测轮廓，mask必须是二值图像
    bounding_boxes = [cv2.boundingRect(cnt) for cnt in in contours]  # 返回下，x,y,w,h,(x,y)左上角的坐标
    bbox_area = 0
    final_bbox = []
    for bbox in bounding_boxes:
        [x,y,w,h] = bbox
        if w*h > bbox_area:
            bbox_area = w*h
            final_bbox = bbox
    return final_bbox
           
if __name__ == "__main__":
    ratio_range = [0.33,0.66]
    composited_dir = '/home/malinjie/composed_lianxi/composited'
    comp_dir = os.path.join(composited_dir, 'composite')
    bg_dir = os.path.join(composited_dir, 'background')
    mask_dir = os.path.join(composited_dir, 'mask')
    if not os.path.exists(comp_dir):
        os.makedirs(comp_dir)
    if not os.path.exists(bg_dir):
        os.makedirs(bg_dir)
    if not os.path.exists(mask_dir):
        os.makedirs(mask_dir)

    figure_files = os.listdir('/home/malinjie/composed_lianxi/img')
    figure_mask_files = os.listdir('/home/malinjie/composed_lianxi/mask')

    figure_files = [os.path.join('/home/malinjie/composed_lianxi/img', fn) for fn in figure_files]  # 将图片路径保存为一个大列表
    figure_mask_files = [os.path.join('/home/malinjie/composed_lianxi/mask', fn) for fn in figure_mask_files]

    for root ,dirs, files in os.walk('/home/malinjie/composed_lianxi/background'):
        for d in dirs:
            path = os.path.join(root, d)
            files = os.listdir(path)
            for bg_fn in files:
                name = os.path.join(path, bg_fn)
                bg = cv2.imread(name)
                bg = cv2.resize(bg, dsize=(1920,1080))      # 读取背景
                bg_id = bg_fn.split('.')[0]
                # random select 5 figures 
                figures_mask_fn = np.random.choice(figure_mask_files, 5,, replace=False)  # 一张背景对应5张mask
                for fg_mask_fn in figures_mask_fn:
                    fg_mask = cv2.imread(fg_mask_fn)[...,0:1]        # 随机读取5张mask
                    fg_id = fg_mask_fn.split('/')[-1].split('.')[0]
                    fg_fn = fg_id + '.jpg'
                    fg = cv2.imread(os.path.join('/home/malinjie/composed_lianxi/img', fg_fn))  # 读取对应5张img
                     # print(fg_mask.shape)
                     bbox = find_boundRect(fg_mask)
                     assert fg.shape[:2] == fg_mask.shape[:2], "Mismatch:\t{}\t{},{}".format(fg_fn, str(fg.shape), str(fg_mask.shape))
                     fg_mask = fg_mask / 255.
                     fg = fg * fg_mask
                     # crop
                     cropped_fg_mask =fg_mask[bbox[1]:bbox[1]+bbox[3],bbox[0]:bbox[0]+bbox[2]]
                     cropped_fg = fg[bbox[1]:bbox[1]+bbox[3], bbox[0]:bbox[0]+bbox[2]]
                
                     # resized
                     h,w,_ = bg.shape
                     fg_h = np.random.randint(int(ratio_range[0]*h), int(ratio_range[1] *h))
                     fg_w = int(cropped_fg.shape[1] / cropped_fg.shape[0] * fg_h)
                     resized_fg = cv2.resize(cropped_fg, dsize=(fg_w, fg_h))
                     resized_mask = cv2.resize(cropped_fg_mask, dsize=(fg_w, fg_h))[..., np.newaxis].repeat(3,axis=2)
                     # cv2.imshow("cropped_fg", cropped_fg)
                     # cv2.imshow("resized_fg", resized_fg)
                     # paste
                     pos_w = np.random.randint(0, w-int(0.5*resized_fg.shape[1]))
                     pos_h = np.random.randint(ratio_range[0]*h, h - int(0.5*resized_fg.shape[0]))
                     paste_max_h = min(h, pos_h+resized_fg.shape[0])
                     paste_max_w = min(w, pos_w+resized_fg.shape[1])
                     resized_fg = resized_fg[:paste_max_h-pos_h, :paste_max_w-pos_w]
                     resized_mask = resized_mask[:paste_max_h-pos_h, :paste_max_w-pos_w]

                     new_comp = np.copy(bg)
                     new_comp[pos_h:paste_max_h, pos_w:paste_max_w] = new_comp[pos_h:paste_max_h, pos_w:paste_max_w]*(1-resized_mask) +  resized_fg* resized_mask
                     new_mask = np.zeros_like(bg)
                     new_mask[pos_h:paste_max_h, pos_w:paste_max_w] = resized_mask
                     new_mask = (new_mask*255).astype(np.uint8)

                     # out = np.concatenate([bg, new_mask], axis=1)
                     # cv2.imshow("out", out)
                     # cv2.waitKey(0)
                
                     # write results
                     final_comp_dir = os.path.join(comp_dir, d)
                     final_mask_dir = os.path.join(mask_dir, d)
                     final_bg_dir = os.path.join(bg_dir, d)

                    if not os.path.exists(final_bg_dir): os.makedirs(final_bg_dir)
                    if not os.path.exists(final_mask_dir): os.makedirs(final_mask_dir)
                    if not os.path.exists(final_comp_dir): os.makedirs(final_comp_dir)

                    cv2.imwrite(os.path.join(final_comp_dir, bg_id+'_' +fg_id+'.jpg'), new_comp)
                    cv2.imwrite(os.path.join(final_bg_dir, bg_id +'_'+fg_id+'.jpg'), bg)
                    cv2.imwrite(os.path.join(final_mask_dir, bg_id+'_'+fg_id+'.jpg'), new_mask)
               print("Generating the synthetic images:\t{}".format(bg_fn))