Copy inference/scripts/pick_cmp_rows.py

inference/scripts/pick_cmp_rows.py

@@ -0,0 +1,64 @@
+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')