inference/scripts/syn_data.py

import os
import random
import os.path as osp
import numpy as np
from pathlib import Path
import shutil
import sys
sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))

from PIL import Image
from  tqdm import tqdm
import click
import cv2

from utils.cv import mask2rle, rle2mask, mask_xyxy
from utils.io_utils import load_exec_list, find_all_files_recursive, find_all_files_with_name, pil_ensure_rgb, imglist2imgrid, imread, imwrite, json2dict, save_tmp_img, dict2json
from utils.sampler import NameSampler
from utils.visualize import visualize_segs, visualize_segs_with_labels, visualize_pos_keypoints
from live2d.scrap_model import Live2DScrapModel, VALID_BODY_PARTS_V1, VALID_BODY_PARTS_V2, compose_mask_from_drawables, init_drawable_visible_map, load_detected_character, load_pos_estimation


@click.group()
def cli():
    """live2d data processing related scripts.
    """


def get_unique_render_lst(exec_list):
    unique_lst = []
    processed_models = set()
    unique_src_to_models = dict()
    for p in tqdm(exec_list):
        modeld = osp.dirname(p)
        if modeld not in processed_models:
            processed_models.add(modeld)
        else:
            continue
        plist = sub_render_parts([p])
        mlist = [Live2DScrapModel(p, skip_load=True) for p in plist]
        for m in mlist:
            m.init_drawables()
        unique_mlist = [mlist[4]]
        for m in mlist:
            is_unique = True
            mklist = list(m.did2drawable.keys())
            mklist.sort()
            for um in unique_mlist:
                umklist = list(um.did2drawable.keys())
                umklist.sort()
                if mklist == umklist:
                    srcp = um.directory
                    is_unique = False
                    break

            tgtp = m.directory
            if is_unique:
                unique_mlist.append(m)
                srcp = m.directory
            if srcp not in unique_src_to_models:
                unique_src_to_models[srcp] = []
            unique_src_to_models[srcp].append(tgtp)
                
        unique_mlist = [m.directory for m in unique_mlist]
        unique_lst += unique_mlist

    return unique_lst, unique_src_to_models


@cli.command('get_tgt_list')
@click.option('--src_dir')
@click.option('--savep', default=None)
def get_tgt_list(src_dir, savep):
    if savep is None:
        savep = osp.join('workspace/datasets', osp.basename(src_dir) + '.txt')

    valid_list = []
    for f in find_all_files_recursive(src_dir, ext={'.json'}):
        tgtf = f.rstrip('.json') + '.png'
        if osp.exists(tgtf):
            valid_list.append(tgtf)
    print(len(valid_list))

    with open(savep, 'w', encoding='utf8') as f:
        f.write('\n'.join(valid_list))


@cli.command('get_png_list')
@click.option('--src_dir')
@click.option('--savep', default=None)
def get_png_list(src_dir, savep):
    if savep is None:
        savep = osp.join('workspace/datasets', osp.basename(src_dir) + '.txt')

    valid_list = []
    for f in find_all_files_recursive(src_dir, ext={'.png'}):
        valid_list.append(f)
    print(len(valid_list))

    with open(savep, 'w', encoding='utf8') as f:
        f.write('\n'.join(valid_list))


@cli.command('check_unique_rst')
@click.option('--exec_list')
@click.option('--savep', default=None)
def check_unique_rst(exec_list, savep):
    if savep is None:
        savep = exec_list
    exec_list = load_exec_list(exec_list)
    exec_list, unique_src_to_models = get_unique_render_lst(exec_list)
    print(len(exec_list))

    with open(savep, 'w', encoding='utf8') as f:
        f.write('\n'.join(exec_list))
    dict2json(unique_src_to_models, savep + '.json')


@cli.command('compress_live2d')
@click.option('--src_dir')
@click.option('--save_dir')
@click.option('--ext', default='.jxl')
@click.option('--disable_crop', is_flag=True, default=False)
def compress_live2d(src_dir, save_dir, ext, disable_crop):

    src_dir = osp.normpath(src_dir)
    model_final_list = find_all_files_with_name(src_dir, 'final')

    crop = not disable_crop
    if save_dir is None:
        save_dir = src_dir +  f'_{ext}'
        if crop:
            save_dir += '_crop'
    save_dir = osp.normpath(save_dir)
    os.makedirs(save_dir, exist_ok=True)

    ndir_leading = len(src_dir.split(os.path.sep))
    for model_f in tqdm(model_final_list, desc=f'saving to {save_dir}'):
        model_dir = osp.dirname(model_f)
        model_save_dir = model_dir.split(os.path.sep)[ndir_leading:]
        model = Live2DScrapModel(model_dir, crop_to_final=crop, pad_to_square=False)
        model.save_model_to(osp.join(save_dir, *model_save_dir), 
                            crop_to_final=crop, img_ext=ext)


@cli.command('build_live2d_exec_list')
@click.option('--srcd')
@click.option('--save_dir', default=None)
@click.option('--filter_p', default=None)
@click.option('--target_fno', default=-1)
@click.option('--num_chunk', default=-1)
@click.option('--save_name', default='exec_list')
def build_live2d_exec_list(srcd, save_dir, filter_p, target_fno, num_chunk, save_name):

    exec_list = find_all_files_with_name(srcd, name='final', exclude_suffix=True)

    tgt_list = []
    filter_set = set()
    if filter_p is not None:
        filter_set = set(load_exec_list(filter_p))
    for d in exec_list:
        if d in filter_set or osp.dirname(d) in filter_set:
            continue
        dname = osp.basename(osp.dirname(d))
        if target_fno > 0:
            fno = dname.split('-')[-1]
            if not fno.isdigit():
                print(f'{d} is not a valid path')
                continue
            fno = int(fno)
            if fno == target_fno:
                tgt_list.append(d)
        else:
            tgt_list.append(d)

    random.shuffle(tgt_list)
    print(f'num samples: {len(tgt_list)}')

    if save_dir is None:
        save_dir = srcd

    with open(osp.join(save_dir, f'{save_name}.txt'), 'w', encoding='utf8') as f:
        f.write('\n'.join(tgt_list))

    if num_chunk > 0:
        world_size = num_chunk
        for ii in range(world_size):
            t = load_exec_list(tgt_list, ii, world_size=world_size)
            with open(osp.join(save_dir, f'{save_name}{ii}.txt'), 'w', encoding='utf8') as f:
                f.write('\n'.join(t))
            print(f'chunk {ii} num samples: {len(t)}')




@cli.command('render_face_samples')
@click.option('--exec_list')
@click.option('--bg_list')
@click.option('--save_dir')
@click.option('--rank_to_worldsize', default='', type=str)
def render_face_samples(exec_list, bg_list, save_dir, rank_to_worldsize):

    TARGET_FRAME_SIZE = 2048

    from utils.cv import fgbg_hist_matching, quantize_image, random_crop, img_bbox, img_alpha_blending, resize_short_side_to, batch_save_masks, batch_load_masks
    from utils.torch_utils import seed_everything
    from utils.visualize import FACE_LABEL2NAME

    def _compose_face_samples(lmodel: Live2DScrapModel):
        '''
        todo: save complete part
        '''
        face_xyxy = lmodel.face_seg_xyxy
        face_h, face_w = face_xyxy[3] - face_xyxy[1], face_xyxy[2] - face_xyxy[0]
        all_face_labels = list(FACE_LABEL2NAME.keys())
        face_final = lmodel.compose_face_drawables(list(FACE_LABEL2NAME.keys()), xyxy=face_xyxy)
        # save_tmp_img(face_final, 'local_tmp.png')

        part_mask_list = []
        # segmap = np.zeros((face_h, face_w), dtype=np.int32)
        alphas = np.zeros((face_h, face_w), dtype=np.int32)
        for ii in range(1, len(all_face_labels)):
            m  = lmodel.compose_face_drawables(ii, mask_only=True, xyxy=face_xyxy, final_visible_mask=True).astype(np.uint8)
            # save_tmp_img(m, mask2img=True)
            part_mask_list.append(m)

        mask_bg = np.bitwise_not(np.bitwise_or.reduce(np.stack(part_mask_list).astype(bool), axis=0))
        part_mask_list.insert(0, mask_bg.astype(np.uint8))
        
        nose_detected, mouth_detected = lmodel.face_part_detected([10, 11])
        tp = osp.join(lmodel.directory, 'faceseg_nosemouth.json.gz')
        if osp.exists(tp) and (not nose_detected or not mouth_detected):
            nose_mouth = batch_load_masks(tp)
            if not nose_detected:
                part_mask_list[10] = nose_mouth[0]
                part_mask_list[1][np.where(nose_mouth[0] > 0)] = 0
            if not mouth_detected:
                part_mask_list[11] = nose_mouth[1]
                part_mask_list[1][np.where(nose_mouth[1] > 0)] = 0

        bx, by, bw, bh = cv2.boundingRect(cv2.findNonZero(part_mask_list[0].astype(np.uint8)))
        by2 = by + bh
        bx2 = bw + bx

        # DONT DELETE THESE!!!!
        
        # depth_lower = 100000
        # depth_upper = -1
        # for d_id, drawable in enumerate(lmodel.drawables):
        #     if drawable.area < 1 or not drawable.face_part_id == 1:
        #         continue

        #     dx, dy, dw, dh = drawable.get_bbox(xyxy=face_xyxy)
        #     dx2 = dx + dw
        #     dy2 = dy + dh

        #     # check if hair drawable is actually background
        #     if drawable.face_part_id == 17:
        #         if drawable.face_part_stats['ioa'][0] > 0.7 and drawable.face_part_stats['ioa'][17] < 0.3:
        #             drawable.face_part_id = None

        #     if drawable.face_part_id == 1 and dw / bw > 0.7 and dh > bw > 0.7:
        #         if drawable.draw_order < depth_lower:
        #             depth_lower = drawable.draw_order
        #         if drawable.draw_order > depth_upper:
        #             depth_upper = drawable.draw_order

        # depth_buffer = np.zeros((face_h, face_w), dtype=np.uint8)
        # base_depth = 1
        # mask = np.zeros_like(depth_buffer, dtype=bool)
        # valid_face_ids = set(range(1, 19))
        # for d in lmodel.drawables:
        #     if d.area < 1 or d.face_part_id not in valid_face_ids:
        #         continue
        #     if np.any(d.bitwise_and(mask, face_xyxy)):
        #         base_depth += 1
        #     m = d.get_full_mask(xyxy=face_xyxy)
        #     mask |= m 
        #     d.depth = base_depth
        #     depth_buffer[np.where(m)] = base_depth
        
        # depth = (depth_buffer / np.max(depth_buffer) * 255).astype(np.uint8)
        # save_tmp_img(depth)


        # base_face_mask = compose_from_drawables([d for d in lmodel.drawables if \
        #                                          drawable.draw_order >= depth_lower and drawable.draw_order > depth_upper])
        # for drawable in lmodel.drawables:
        #     if drawable.draw_order < depth_lower or drawable.draw_order > depth_upper:
        #         continue

        # segmap = segmap.astype(np.uint8)
        # lmodel.compose_face_drawables([4, 5], output_type='pil').save('local_tst.png')
        # save_tmp_img(face_final)
        # save_tmp_img(segmap == 1, mask2img=True)
        # save_tmp_img(segmap == 4, mask2img=True)

        return True, part_mask_list, face_final
    
    os.makedirs(save_dir, exist_ok=True)

    seed_everything(42)

    hist_match_prob = 0.2
    quantize_prob = 0.25
    color_correction_sampler = NameSampler({'hist_match': hist_match_prob, 'quantize': quantize_prob})
    
    if exec_list.endswith('.json'):
        new_exec_list = []
        exec_list = json2dict(exec_list)
        for k, vs in exec_list.items():
            for v in vs:
                new_exec_list.append({v: k})
        exec_list = new_exec_list
        pass
    
    exec_list = load_exec_list(exec_list, rank_to_worldsize=rank_to_worldsize)
    bg_list = load_exec_list(bg_list)

    VALID_FACE_SET = set(range(19))
    for ii, p in enumerate(tqdm(exec_list[0:])):
        try:
            face_parsingp = None
            if isinstance(p, dict):
                for k, v in p.items():
                    p = k
                    face_parsingp = osp.join(v, 'face_parsing.json')
            lmodel = Live2DScrapModel(p)
            model_dir = lmodel.directory

            if face_parsingp is None:
                face_parsingp = osp.join(model_dir, 'face_parsing.json')
                if not osp.exists(face_parsingp):
                    face_parsingp = '-'.join(model_dir.split('-')[:-1]) + '-4'
                    face_parsingp = osp.join(face_parsingp, 'face_parsing.json')
            if not osp.exists(face_parsingp):
                print(f"skip {p} due to face parsing result not found")
                continue
            
            lmodel.load_face_parsing(face_parsingp)
            face_drawables = [d for d in lmodel.drawables if d.face_part_id in VALID_FACE_SET]
            init_drawable_visible_map(face_drawables)

            is_valid, labels, face_final = _compose_face_samples(lmodel,)
            mask_list = labels
            if not is_valid:
                continue
            # save_tmp_img(labels[0], mask2img=True)

            bgp = random.choice(bg_list)
            fh, fw = face_final.shape[:2]
            bg = imread(bgp)
            bgh, bgw = bg.shape[:2]
            target_bg_size = min(bgh, bgw, TARGET_FRAME_SIZE)
            fsize = max(fh, fw)
            if fsize * 2 < target_bg_size:
                target_bg_size = random.randint(fsize * 2, target_bg_size)
            bg = resize_short_side_to(bg, target_bg_size)
            bg = random_crop(imread(bgp), (fh, fw))
            # save_tmp_img(bg)
            
            color_correct = color_correction_sampler.sample()
            if color_correct == 'hist_match':
                fgbg_hist_matching([face_final], bg)

            face_wbg = img_alpha_blending([bg, face_final])
            
            if color_correct == 'quantize':
                mask = face_final[..., -1] > 35
                # cv2.imshow("mask", mask)
                face_wbg[..., :3] = quantize_image(face_wbg[..., :3], random.choice([12, 16, 32]), 'kmeans', mask=mask)[0]
            
            d = osp.abspath(model_dir).replace('\\', '/').rstrip('/').replace('.', '_DOT_')
            d1 = d.split('/')[-1]
            d2 = d.split('/')[-3]
            savename = d2 + '____' + d1
            savep = osp.join(save_dir, savename)
            # save_tmp_img(face_wbg)
            imwrite(savep, face_wbg, quality=97, ext='.jpg')
            batch_save_masks(mask_list, savep + '.json', compress='gzip')
            # print(f'finished {savep}')
        except Exception as e:
            # raise
            print(f'Failed to process {p}: {e}')


@cli.command('get_tgt_list')
@click.option('--src_dir')
@click.option('--savep', default=None)
def get_tgt_list(src_dir, savep):
    if savep is None:
        savep = osp.join('workspace/datasets', osp.basename(src_dir) + '.txt')

    valid_list = []
    for f in find_all_files_recursive(src_dir, ext={'.json'}):
        tgtf = osp.splitext(f)[0] + '.png'
        if osp.exists(tgtf):
            valid_list.append(tgtf)
    print('valid samples: ', len(valid_list))

    with open(savep, 'w', encoding='utf8') as f:
        f.write('\n'.join(valid_list))
    print(f'exec_list saved to {savep}')

@cli.command('render_body_samples')
@click.option('--exec_list')
@click.option('--bg_list')
@click.option('--mask_name', default=None)
@click.option('--save_dir', default='')
@click.option('--rank_to_worldsize', default='', type=str)
@click.option('--save_suffix', default='.png', type=str)
def render_body_samples(exec_list, bg_list, mask_name, save_dir, rank_to_worldsize, save_suffix):

    from live2d.scrap_model import animal_ear_detected, Drawable, ImageProcessor, compose_from_drawables, VALID_BODY_PARTS_V3
    from utils.cv import fgbg_hist_matching, quantize_image, random_crop, rle2mask, mask2rle, img_alpha_blending, resize_short_side_to, batch_save_masks, batch_load_masks
    from utils.torch_utils import seed_everything

    seed_everything(42)

    hist_match_prob = 0.35
    # quantize_prob = 0.25
    color_correction_sampler = NameSampler({'hist_match': hist_match_prob, 'quantize': 0.})
    
    exec_list = load_exec_list(exec_list, rank_to_worldsize=rank_to_worldsize)
    bg_list = load_exec_list(bg_list)

    tagcluster_bodypart = json2dict('common/assets/tagcluster_bodypart_v2.json')
    tag2generaltag = {}
    for general_tag, tlist in tagcluster_bodypart.items():
        for t in tlist:
            if t in tag2generaltag and tag2generaltag[t] != general_tag:
                print(f'conflict tag def: {t} - {general_tag}, ' + tag2generaltag[t])
            tag2generaltag[t] = general_tag

    if save_dir != '':
        os.makedirs(save_dir, exist_ok=True)
    render_sample = save_dir != ''

    MAX_TGT_SIZE = 1280
    target_tag_list = VALID_BODY_PARTS_V3 + ['head']

    invalid_lst: list[int] = [2094, 1389, 627, 477, 280, 480]


    for ii, p in enumerate(tqdm(exec_list)):
        try:
            lmodel = Live2DScrapModel(p)
            load_success = lmodel.load_body_parsing(mask_name)
            if not load_success:
                print(f'failed to load body parsing, skip: {p}')
                continue

            metadata = lmodel._body_parsing['metadata']
            if metadata is None:
                metadata = {}
            is_valid = metadata.get('is_valid', True)
            is_incomplete = metadata.get('is_incomplete', False)
            is_cleaned = metadata.get('cleaned', False)
            tag_valid = metadata.get('tag_valid', {})
            object_valid = True
            foot_valid = True

            if not is_valid:
                continue

            # if is_incomplete:
            #     continue

            # keep_bg = random.random() < 0.3
            keep_bg = False

            if not is_valid:
                continue
            
            valid_drawables: list[Drawable] = []
            body_drawables: list[Drawable] = []
            h, w = lmodel.size()
            x_min, x_max, y_min, y_max = w, 0, h, 0
            for d in lmodel.drawables:
                d.get_img()
                if d.area < 1:
                    continue
                if not keep_bg and d.body_part_tag not in target_tag_list:
                    continue
                valid_drawables.append(d)
                if d.body_part_tag in target_tag_list:
                    body_drawables.append(d)
                dxyxy = d.xyxy
                x_min = min(x_min, dxyxy[0])
                x_max = max(x_max, dxyxy[2])
                y_min = min(y_min, dxyxy[1])
                y_max = max(y_max, dxyxy[3])

            if keep_bg:
                x_min = y_min = 0
                x_max = w
                y_max = h

            ch, cw = y_max - y_min, x_max - x_min
            scale = min(MAX_TGT_SIZE / max(ch, cw), 1)
            nh, nw = ch, cw
            if scale < 1:
                nh = int(round(nh * scale))
                nw = int(round(nw * scale))
            new_processor = ImageProcessor(target_frame_size=[nw, nh], crop_bbox=[x_min, y_min, x_max, y_max], pad_to_square=False)
            lmodel.final = new_processor(lmodel.final, update_coords_modifiers=True)
            lmodel.final_bbox = [
                new_processor.crop_bbox[0] + x_min,
                new_processor.crop_bbox[1] + y_min,
                new_processor.crop_bbox[0] + x_max,
                new_processor.crop_bbox[1] + y_max
            ]
            for d in valid_drawables:
                d.set_img_processor(new_processor)
                d._final_size = [nh, nw]
                d.load_img(force_reload=True, img=d.img)
            
            h, w = lmodel.size()
            depth_buffer = np.zeros((h, w), dtype=np.uint16)
            base_depth = 1
            init_drawable_visible_map(valid_drawables)
            # part_mask_list, body_final = _compose_body_samples(lmodel)
            
            part_mask_list = []
            if not keep_bg:
                body_final = lmodel.compose_bodypart_drawables(target_tag_list)
            else:
                body_final = compose_from_drawables(valid_drawables)

            for tag in target_tag_list:
                m  = lmodel.compose_bodypart_drawables(tag, mask_only=True, final_visible_mask=True).astype(np.uint8)
                part_mask_list.append(m)
            
            mask = np.zeros((h, w), dtype=bool)
            for d in body_drawables:
                m = d.get_full_mask()
                if np.any(d.bitwise_and(mask, [0, 0, w, h])):
                    base_depth += 1
                    mask = m
                else:
                    mask |= m                
                d.depth = base_depth
                depth_buffer[np.where(m)] = base_depth
            
            depth_dtype = np.uint8
            if base_depth > 255:
                depth_dtype = np.uint16
            depth_buffer = depth_buffer.astype(depth_dtype)

            d = osp.abspath(lmodel.directory).replace('\\', '/').rstrip('/').replace('.', '_DOT_')
            d1 = d.split('/')[-1]
            d2 = d.split('/')[-3]
            savename = d2 + '____' + d1
            savep = osp.join(save_dir, savename)

            masks = part_mask_list
            foot_msk_idx = target_tag_list.index('footwear')
            object_msk_idx = target_tag_list.index('objects')
            leg_msk_idx = target_tag_list.index('legwear')
            masks[leg_msk_idx] = masks[leg_msk_idx] | masks[foot_msk_idx]
            px = py = 0
            
            final_img = body_final

            bgp = random.choice(bg_list)
            fh, fw = final_img.shape[:2]
            bg = imread(bgp)
            fsize = min(max(h, w), MAX_TGT_SIZE)
            fsze_max = int(round(fsize * 1.5))
            target_bg_size = random.randint(fsize, fsze_max)
            bg = resize_short_side_to(bg, target_bg_size)

            target_bg_w = target_bg_h = target_bg_size
            if fh > fw:
                target_bg_w = random.randint(fw, target_bg_size)
            elif fw > fh:
                target_bg_h = random.randint(fh, target_bg_size)

            bg = random_crop(imread(bgp), (target_bg_h, target_bg_w))
            
            px = py = 0
            if fh != target_bg_h or fw != target_bg_w:
                if fh != target_bg_h:
                    py = random.randint(0, target_bg_h - fh)
                if fw != target_bg_w:
                    px = random.randint(0, target_bg_w - fw)
                blank_final = np.zeros((target_bg_h, target_bg_w, 4), np.uint8)
                blank_final[py: py + fh, px: px + fw] = final_img
                final_img = blank_final

                depth_blank = np.zeros((target_bg_h, target_bg_w), dtype=depth_dtype)
                depth_blank[py: py + fh, px: px + fw] = depth_buffer
                depth_buffer = depth_blank

                for mi, m in enumerate(masks):
                    blank = np.zeros((target_bg_h, target_bg_w), bool)
                    blank[py: py + fh, px: px + fw] = m
                    masks[mi] = blank
            fh, fw = final_img.shape[:2]

            color_correct = color_correction_sampler.sample()

            if color_correct == 'hist_match':
                fgbg_hist_matching([final_img], bg, fg_only=True)

            wbg = img_alpha_blending([bg, final_img])
            wbg[..., -1] = final_img[..., -1]


            fh, fw = wbg.shape[:2]

            # save_tmp_img(visualize_segs_with_labels(masks, wbg[..., :3], tag_list=target_tag_list, image_weight=0.1))
            imwrite(savep, wbg, quality=100, ext=save_suffix)
            imwrite(savep + '_depth', depth_buffer, quality=100, ext='.png')
            
            mask_meta_list = [{} for _ in range(len(target_tag_list))] # dont use [{}] * len
            mask_meta_list[foot_msk_idx]['is_valid'] = foot_valid
            mask_meta_list[object_msk_idx]['is_valid'] = object_valid
            batch_save_masks(masks, savep + '.json', mask_meta_list=mask_meta_list)
            del masks
            # del wbg
            del depth_buffer

            sample_ann = {'cleaned': is_cleaned, 'is_incomplete': is_incomplete, 'tag_info': {k: {'valid': True, 'exists': False} for k in VALID_BODY_PARTS_V2}, 'final_size': wbg.shape[:2]}
            tag_info = sample_ann['tag_info']
            # tag_info['footwear']['valid'] = foot_valid
            # tag_info['objects']['valid'] = object_valid

            for ii, tag in enumerate(target_tag_list):
                # if tag == 'footwear' and not foot_valid:
                #     continue
                # if tag == 'objects' and not object_valid:
                #     continue
                if tag == 'head':
                    drawables = lmodel.get_body_part_drawables(['face', 'irides', 'eyebrow', 'eyewhite', 'eyelash', 'eyewear', 'ears', 'nose', 'mouth'])
                else:
                    drawables = lmodel.get_body_part_drawables(tag)
                # if tag == 'legwear':
                #     drawables += lmodel.get_body_part_drawables('footwear')
                drawables = [d for d in drawables if d.area >= 1]
                if len(drawables) == 0:
                    continue
                init_drawable_visible_map(drawables)
                x_min, x_max, y_min, y_max = fw, 0, fh, 0
                for d in drawables:
                    dxyxy = d.xyxy
                    x_min = min(x_min, dxyxy[0])
                    x_max = max(x_max, dxyxy[2])
                    y_min = min(y_min, dxyxy[1])
                    y_max = max(y_max, dxyxy[3])
                
                xyxy = [x_min, y_min, x_max, y_max]
                dh, dw = y_max - y_min, x_max - x_min
                part_final = compose_from_drawables(drawables, xyxy=xyxy)
                imwrite(savep + f'_{tag}', part_final, quality=100, ext='.png')
                
                depth_buffer = np.zeros((dh, dw), dtype=depth_dtype)
                for d in drawables:
                    dxyxy = d.xyxy
                    m = d.final_visible_mask
                    depth_buffer[dxyxy[1] - y_min: dxyxy[3] - y_min, dxyxy[0] - x_min: dxyxy[2] - x_min][np.where(m)] = d.depth
                
                xyxy = [x_min + px, y_min + py, x_max + px, y_max + py]
                imwrite(savep + f'_{tag}_depth', depth_buffer, quality=100, ext='.png')
                if tag not in tag_info:
                    tag_info[tag] = {}
                tag_info[tag]['exists'] = True
                tag_info[tag]['xyxy'] = xyxy

                blank = np.zeros_like(wbg)
                blank[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]] = part_final
                # save_tmp_img(wbg)
                # save_tmp_img(img_alpha_blending([wbg, blank]))
                # pass
                
            dict2json(sample_ann, savep + '_ann.json')

        except Exception as e:
            raise
            print(f'Failed to process {p}: {e}')





if __name__ == '__main__':
    cli()