import os

default_n_threads = 8
os.environ['OPENBLAS_NUM_THREADS'] = f"{default_n_threads}"
os.environ['MKL_NUM_THREADS'] = f"{default_n_threads}"
os.environ['OMP_NUM_THREADS'] = f"{default_n_threads}"

from utils.cv import *
from utils.io_utils import *
import shutil

lst = ['photo_2026-01-03_01-03-02', 'Generated Image January 03, 2026 - 12_47AM', 'Generated Image January 03, 2026 - 12_00AM']

dirs = ['tmp/cmp_part_extr/ours', 'workspace/datasets/testcaseall_output_woattn']
heads = ['ours', 'woattn']

save_dir = 'tmp/woattn_cmp'

for imgn in lst:
    for ii, d in enumerate(dirs):
        src_dir = osp.join(d, imgn)
        from talking_head.preprocess import further_extr
        further_extr(src_dir, rotate=False)

        saved = save_dir
        os.makedirs(saved, exist_ok=True)

        if ii == 0:
            src_img = np.array(Image.open(osp.join(src_dir, 'src_img.png')))
            sz = src_img.shape[:2]
            xyxy = np.array(cv2.boundingRect(cv2.findNonZero(src_img[..., -1])))
            xyxy[[2, 3]] += xyxy[[0, 1]]
            img = src_img[xyxy[1]: xyxy[3], xyxy[0]: xyxy[2]].copy()
            
            save_tmp_img(img, osp.join(saved, imgn + '_input.png'))

        img_list = []


        src_infop = osp.join(osp.join(src_dir, 'optimized'), 'info.json')
        infos = json2dict(src_infop)
        load_img_depth(osp.join(src_dir, 'optimized'), infos, pad=0)

        flist = []
        for k, v in infos['parts'].items():
            # depth_median = v['depth'][v['img'][..., -1] > 127]
            # dm = np.median(depth_median)
            # v['depth_median'] = dm
            if k =='armlf':
                v['depth_median'] = -1.
            if k == 'bottomwear':
                continue
            v.pop('depth')
            save_tmp_img(v['img'])
            flist.append(v)

        flist.sort(key = lambda x: x['depth_median'], reverse=True)
        # for p in
                
        img = img_alpha_blending(
            flist, final_size=sz, premultiplied=False)
        
        img = img
        save_tmp_img(img, osp.join(saved, imgn + '_' + heads[ii]) + '.png')