import os.path as osp import os from pathlib import Path from typing import Union, List, Dict import json from collections.abc import MutableMapping import gzip from functools import lru_cache import yaml import numpy as np from PIL import Image import cv2 # numpy 2.x compatible — np.bool, np.bool8, np.float_, np.int_, np.uint were removed NP_BOOL_TYPES = (np.bool_,) NP_FLOAT_TYPES = (np.float16, np.float32, np.float64) NP_INT_TYPES = (np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64) class NumpyEncoder(json.JSONEncoder): def default(self, obj): if isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, np.ScalarType): if isinstance(obj, NP_BOOL_TYPES): return bool(obj) elif isinstance(obj, NP_FLOAT_TYPES): return float(obj) elif isinstance(obj, NP_INT_TYPES): return int(obj) return json.JSONEncoder.default(self, obj) def json2dict(json_path: str): plower = json_path.lower() if plower.endswith('.gz'): with gzip.open(json_path, 'rt', encoding='utf8') as f: metadata = json.load(f) return metadata if plower.endswith('.yaml'): with open(json_path, 'r') as file: metadata = yaml.load(file, yaml.CSafeLoader) return metadata with open(json_path, 'r', encoding='utf8') as f: metadata = json.loads(f.read()) return metadata def serialize_np(obj): if isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, np.ScalarType): if isinstance(obj, NP_BOOL_TYPES): return bool(obj) elif isinstance(obj, NP_FLOAT_TYPES): return float(obj) elif isinstance(obj, NP_INT_TYPES): return int(obj) return obj def json_dump_nested_obj(obj, **kwargs): def _default(obj): if isinstance(obj, (np.ndarray, np.ScalarType)): return serialize_np(obj) return obj.__dict__ return json.dumps(obj, default=lambda o: _default(o), ensure_ascii=False, **kwargs) def dict2json(adict: dict, json_path: str, compress=None): if compress is None: with open(json_path, "w", encoding="utf-8") as f: f.write(json.dumps(adict, ensure_ascii=False, cls=NumpyEncoder)) elif compress == 'gzip': if not json_path.endswith('.gz'): json_path += '.gz' with gzip.open(json_path, 'wt', encoding="utf8") as zipfile: json.dump(adict, zipfile, ensure_ascii=False, cls=NumpyEncoder) else: raise Exception(f'Invalid compression: {compress}') IMG_EXT = ['.bmp', '.jpg', '.png', '.jpeg', '.webp', '.jxl'] def find_all_imgs(img_dir, abs_path=False, sort=False): imglist = [] dir_list = os.listdir(img_dir) for filename in dir_list: if filename.startswith('.'): continue file_suffix = Path(filename).suffix if file_suffix.lower() not in IMG_EXT: continue if abs_path: imglist.append(osp.join(img_dir, filename)) else: imglist.append(filename) if sort: imglist.sort() return imglist def get_last_modified_file(file_prefix, exts, ext_fallback=None): ''' get last modified file from files sharing same prefix ''' latest_time = -1 latest_f = None for ext in exts: tmp_p = file_prefix + ext if osp.exists(tmp_p) and osp.getmtime(tmp_p) > latest_time: latest_time = osp.getmtime(tmp_p) latest_f = tmp_p if latest_f is None: if ext_fallback is not None: latest_f = file_prefix + ext_fallback else: latest_f = file_prefix + exts[0] return latest_f def find_all_files_recursive(tgt_dir: Union[List, str], ext: Union[List, set], exclude_dirs=None): if isinstance(tgt_dir, str): tgt_dir = [tgt_dir] if exclude_dirs is None: exclude_dirs = set() filelst = [] for d in tgt_dir: for root, _, files in os.walk(d): if osp.basename(root) in exclude_dirs: continue for f in files: if Path(f).suffix.lower() in ext: filelst.append(osp.join(root, f)) return filelst def find_all_files_with_name(tgt_dir: Union[List, str], name, exclude_dirs=None, exclude_suffix=True): if isinstance(tgt_dir, str): tgt_dir = [tgt_dir] if exclude_dirs is None: exclude_dirs = set() filelst = [] for d in tgt_dir: for root, _, files in os.walk(d): if osp.basename(root) in exclude_dirs: continue for f in files: fn = osp.basename(f) if exclude_suffix: fn = osp.splitext(fn)[0] if fn == name: filelst.append(osp.join(root, f)) return filelst def find_all_imgs_recursive(tgt_dir, exclude_dirs=None): return find_all_files_recursive(tgt_dir, IMG_EXT, exclude_dirs) VIDEO_EXT = {'.mp4', '.gif', '.webm', '.avif', '.mkv'} def find_all_videos_recursive(tgt_dir,exclude_dirs=None): return find_all_files_recursive(tgt_dir, VIDEO_EXT, exclude_dirs) def load_exec_list(exec_list, rank=None, world_size=None, check_exist=False, to_imgs=False, rank_to_worldsize=None): ''' split exec_list by rank and world_size if available ''' if rank_to_worldsize is not None and rank_to_worldsize != '' and rank_to_worldsize != '-': rank, world_size = rank_to_worldsize.split('-') rank = int(rank) world_size = int(world_size) if isinstance(exec_list, str): if osp.exists(exec_list): if exec_list.endswith('.json') or exec_list.endswith('.json.gz'): exec_list = json2dict(exec_list) elif exec_list.endswith('.txt'): with open(exec_list, 'r', encoding='utf8') as f: exec_list = f.read().split('\n') else: exec_list = [exec_list] else: exec_list = exec_list.split(',') else: exec_list = list(exec_list) if rank is not None and world_size is not None: nexec = len(exec_list) // world_size nstart = nexec * rank if rank == world_size - 1: exec_list = exec_list[nstart:] else: exec_list = exec_list[nstart:nstart+nexec] if to_imgs: _exec_list = [] for p in exec_list: if osp.isdir(p): _exec_list += find_all_imgs(p, sort=True, abs_path=True) else: _exec_list.append(p) exec_list = _exec_list if check_exist: nlist = [] for p in exec_list: if osp.exists(p): nlist.append(p) return nlist else: return exec_list def get_rank(): if 'RANK' in os.environ: # print('worksize: ', os.environ['WORLD_SIZE']) return int(os.environ['RANK']), int(os.environ['WORLD_SIZE']) return None, None def load_image(imgp: str, mode="RGB", output_type='numpy'): """ return RGB image as output_type """ img = Image.open(imgp).convert(mode) if output_type == 'numpy': img = np.array(img) if len(img.shape) == 2: img = img[..., None] return img def flatten_dict(dictionary, parent_key='', separator='_'): items = [] parent_key = str(parent_key) for key, value in dictionary.items(): new_key = parent_key + separator + str(key) if parent_key else str(key) if isinstance(value, MutableMapping): items.extend(flatten_dict(value, new_key, separator=separator).items()) else: items.append((new_key, value)) return dict(items) def imglist2imgrid(imglist, cols=4, output_type='numpy', fix_size=None): if isinstance(fix_size, int): fix_size = (fix_size, fix_size) current_row = [] grid = [] grid.append(current_row) for ii, img in enumerate(imglist): if isinstance(img, Image.Image): img = np.array(img) if fix_size is not None: if fix_size[0] != img.shape[0] or fix_size[1] != img.shape[1]: img = cv2.resize(img, (fix_size[1], fix_size[0]), interpolation=cv2.INTER_AREA) current_row.append(img) if len(current_row) >= cols and ii != len(imglist) - 1: current_row = [] grid.append(current_row) if len(grid) > 1 and len(grid[-1]) < cols: for ii in range(cols - len(grid[-1])): grid[-1].append(np.full_like(grid[-1][-1], fill_value=255)) if len(grid) > 1: for ii, row in enumerate(grid): grid[ii] = np.concatenate(row, axis=1) grid = np.concatenate(grid, axis=0) else: grid = np.concatenate(grid[0], axis=1) if output_type.lower() == 'pil': grid = Image.fromarray(grid) return grid def pil_ensure_rgb(image: Image.Image) -> Image.Image: if isinstance(image, str): image = Image.open(image) is_array = False if isinstance(image, np.ndarray): is_array = True image = Image.fromarray(image) # convert to RGB/RGBA if not already (deals with palette images etc.) if image.mode not in ["RGB", "RGBA"]: image = image.convert("RGBA") if "transparency" in image.info else image.convert("RGB") # convert RGBA to RGB with white background if image.mode == "RGBA": canvas = Image.new("RGBA", image.size, (255, 255, 255)) canvas.alpha_composite(image) image = canvas.convert("RGB") if is_array: image = np.array(image) return image def pil_pad_square(image: Image.Image) -> Image.Image: w, h = image.size # get the largest dimension so we can pad to a square px = max(image.size) # pad to square with white background canvas = Image.new("RGBA", (px, px), (0, 0, 0, 0)) padding = ((px - w) // 2, (px - h) // 2) canvas.paste(image, padding) return canvas, padding def load_facedet_result(srcp: str): preds = json2dict(srcp) for pred in preds: bbox = np.array(pred['bbox'], dtype=np.float32) keypoints = np.array(pred['keypoints'], dtype=np.float32) bbox[-1] = np.round(bbox[-1] / 100) keypoints[:, 2] = np.round(keypoints[:, 2] / 100) pred['bbox'] = bbox pred['keypoints'] = keypoints return preds def intersect_area(xyxy1, xyxy2): l = max(xyxy1[0], xyxy2[0]) r = min(xyxy1[2], xyxy2[2]) t = max(xyxy1[1], xyxy2[1]) b = min(xyxy1[3], xyxy2[3]) if l > r or t > b: return -1 return (r - l) * (b - t) def bbox_iou(xyxy1, xyxy2): i = intersect_area(xyxy1, xyxy2) if i < 0: return i u = (xyxy1[3] - xyxy1[1]) * (xyxy1[2] - xyxy1[0]) + \ (xyxy2[3] - xyxy2[1]) * (xyxy2[2] - xyxy2[0]) - i iou = i / u return iou def imread(imgpath, read_type=cv2.IMREAD_COLOR, max_retry_limit=5, retry_interval=0.1): if not osp.exists(imgpath): return None num_tries = 0 img = Image.open(imgpath) if read_type != cv2.IMREAD_GRAYSCALE: img = img.convert('RGB') img = np.array(img) if read_type == cv2.IMREAD_GRAYSCALE: if img.ndim == 3: if img.shape[-1] == 3: img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) elif img.shape[-1] == 4: img = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY) elif img.shape[-1] == 1: img = img[..., 0] else: raise return img def imwrite(img_path, img, ext='.png', quality=100, jxl_encode_effort=3): # cv2 writing is faster than PIL suffix = Path(img_path).suffix ext = ext.lower() assert ext in IMG_EXT if suffix != '': img_path = img_path.replace(suffix, ext) else: img_path += ext encode_param = None if ext in {'.jpg', '.jpeg'}: encode_param = [cv2.IMWRITE_JPEG_QUALITY, quality] elif ext == '.webp': if quality == 100: quality = 101 encode_param = [cv2.IMWRITE_WEBP_QUALITY, quality] if ext == '.jxl': # jxl_encode_effort: https://github.com/Isotr0py/pillow-jpegxl-plugin/issues/23 # higher values theoretically produce smaller files at the expense of time, 3 seems to strike a balance lossless = quality > 99 # quality=100, lossless=False seems to result in larger file compared with lossless=True Image.fromarray(img).save(img_path, quality=quality, lossless=lossless, effort=jxl_encode_effort) else: if len(img.shape) == 3: if img.shape[-1] == 3: img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) elif img.shape[-1] == 4: img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA) cv2.imencode(ext, img, encode_param)[1].tofile(img_path) def save_tmp_img(img: Union[Image.Image, np.ndarray], savep = 'local_tst.png', mask2img=False): ''' for debug img output ''' if not savep: savep = 'local_tst.png' if mask2img: img = img.astype(np.uint8) * 255 if isinstance(img, np.ndarray): if img.ndim == 3 and img.shape[-1] == 1: img = img[..., 0] img = Image.fromarray(img.astype(np.uint8)) img.save(savep) def bbox2xyxy(box): x1, y1 = box[0], box[1] return x1, y1, x1+box[2], y1+box[3] def bbox_overlap_area(abox, boxb) -> int: ax1, ay1, ax2, ay2 = bbox2xyxy(abox) bx1, by1, bx2, by2 = bbox2xyxy(boxb) ix = min(ax2, bx2) - max(ax1, bx1) iy = min(ay2, by2) - max(ay1, by1) if ix > 0 and iy > 0: return ix * iy else: return 0 def bbox_overlap_xy(abox, boxb): ax1, ay1, ax2, ay2 = bbox2xyxy(abox) bx1, by1, bx2, by2 = bbox2xyxy(boxb) ix = min(ax2, bx2) - max(ax1, bx1) iy = min(ay2, by2) - max(ay1, by1) return ix, iy @lru_cache(maxsize=4) def get_all_segcls(cls_path: str): if cls_path.lower().endswith('.json'): cls_list = list(json2dict(cls_path).keys()) else: with open(cls_path, 'r', encoding='utf8') as f: c = f.read() cls_list = [l.strip() for l in c.split('\n') if l.strip()] return cls_list def imglist_from_dir_or_flist(src): if osp.isdir(src): lst = find_all_imgs(src, sort=True, abs_path=True) else: assert osp.isfile(src) lst = load_exec_list(src) return lst def find_closest_point_from_line2(p0, p1, pts): dist = np.linalg.norm(pts - p0[None, :], axis=1) + np.linalg.norm(pts - p1[None, :], axis=1) return pts[np.argsort(dist)[0]] def cosine_similarity_numpy(v1, v2): """Calculates the cosine similarity between two NumPy vectors.""" v1 = np.array(v1) v2 = np.array(v2) dot_product = np.dot(v1, v2) magnitude_v1 = np.linalg.norm(v1) magnitude_v2 = np.linalg.norm(v2) if magnitude_v1 == 0 or magnitude_v2 == 0: return 0 # Or handle as an error, depending on requirements return dot_product / (magnitude_v1 * magnitude_v2) def xyxy2center(xyxy): xyxy = np.array(xyxy) return np.array(xyxy[[0, 1]] + xyxy[[2, 3]]) / 2 def save_psd(savep, img_list, h, w, pad_to_canvas=False, mode='RGBA', img_key='img'): from psd_tools import PSDImage psd_image = PSDImage.new(mode=mode, size=(h, w), depth=8) for imgd in img_list: img = imgd[img_key] x1 = y1 = 0 if 'xyxy' in imgd: x1, y1, x2, y2 = imgd['xyxy'] if 'xyxy' in imgd and pad_to_canvas: img_padded = np.zeros((h, w, 4), dtype=np.uint8) img_padded[y1: y2, x1: x2] = img x1 = y1 = 0 img = img_padded img = Image.fromarray(img) layer_name = 'undefined' if 'layer_name' in imgd: layer_name = imgd['layer_name'] elif 'tag' in imgd: layer_name = imgd['tag'] psd_image.create_pixel_layer(img, name=layer_name, top=y1, left=x1, opacity=255) psd_image.save(savep) def load_part(srcp: str, rotate=False, pad=0, min_width=64, min_sz=12, depth_min=None, depth_max=None): img = Image.open(srcp).convert('RGBA') srcd = osp.dirname(srcp) tag = osp.splitext(osp.basename(srcp))[0] depthp = osp.join(srcd, tag + '_depth.png') tag_infop = osp.join(srcd, tag + '.json') img = np.array(img) p_test = max(img.shape[:2]) // 10 mask = img[..., -1] > 10 if isinstance(pad, int): pad = [pad] * 4 if osp.exists(tag_infop): rst = json2dict(tag_infop) depth = np.array(Image.open(depthp).convert('L')) depth = np.array(depth, dtype=np.float32) / 255 rst.update({'img': img, 'depth': depth, 'mask': mask, 'tag': tag}) return rst if np.sum(mask[:-p_test, :-p_test]) > 4: if rotate: img = np.rot90(img, 3) mask = np.rot90(mask, 3, ) xyxy = cv2.boundingRect(cv2.findNonZero(mask.astype(np.uint8))) xyxy = np.array(xyxy) h, w = xyxy[2:] xyxy[2] += xyxy[0] xyxy[3] += xyxy[1] p = min_width - w if p > 0: if xyxy[0] >= p: xyxy[0] -= p else: xyxy[2] += p p = min_sz - h if p > 0: if xyxy[1] >= p: xyxy[1] -= p else: xyxy[3] += p depth = np.array(Image.open(depthp).convert('L')) if rotate: depth = np.rot90(depth, 3) x1, y1, x2, y2 = xyxy mask = mask[y1: y2, x1: x2].copy() img = img[y1: y2, x1: x2].copy() depth = depth[y1: y2, x1: x2].copy() pt, pb, pl, pr = pad if pt > 0 or pb > 0 or pl > 0 or pr > 0: img = cv2.copyMakeBorder(img, pt, pb, pl, pr, cv2.BORDER_CONSTANT, value=(0, 0, 0, 0)) depth = cv2.copyMakeBorder(depth, pt, pb, pl, pr, cv2.BORDER_CONSTANT, value=(255)) mask = cv2.copyMakeBorder(mask.astype(np.uint8), pt, pb, pl, pr, cv2.BORDER_CONSTANT, value=(0)) > 0 x1 -= pl y1 -= pt x2 += pr y2 += pb xyxy = [x1, y1, x2, y2] # dmin, dmax = partdict['depth_min'], partdict['depth_max'] depth = np.array(depth, dtype=np.float32) / 255 rst = {'img': img, 'depth': depth, 'xyxy': xyxy, 'mask': mask, 'tag': tag} if depth_max is not None and depth_min is not None: dmax, dmin = depth_max, depth_min depth = depth * (dmax - dmin) + dmin rst.update({'depth': depth, 'depth_min': dmin, 'depth_max': dmax}) return rst else: return None def load_img_depth(srcd, src_info, pad=5, try_crop=False, rotate=False): ''' pad: int or [pt, pb, pl, pr] ''' tag2infos = src_info['parts'] if isinstance(pad, int): pad = [pad] * 4 for t in tag2infos: part_info = tag2infos[t] if 'img' in part_info: continue img = np.array(Image.open(osp.join(srcd, t + '.png')).convert('RGBA')) # if rotate: # img = np.rot90 # if try_crop: # xyxy = cv2.boundingRect(cv2.findNonZero((img[..., -1] > 10).astype(np.uint8))) depth = np.array(Image.open(osp.join(srcd, t + '_depth.png')).convert('L')) if 'depth_max' in part_info: dmax, dmin = part_info['depth_max'], part_info['depth_min'] depth = np.array(depth, dtype=np.float32) / 255 * (dmax - dmin) + dmin else: depth = np.array(depth, dtype=np.float32) / 255 if 'xyxy' in part_info: x1, y1, x2, y2 = part_info['xyxy'] pt, pb, pl, pr = pad if pt > 0 or pb > 0 or pl > 0 or pr > 0: img = cv2.copyMakeBorder(img, pt, pb, pl, pr, cv2.BORDER_CONSTANT, value=(0, 0, 0, 0)) depth = cv2.copyMakeBorder(depth, pt, pb, pl, pr, cv2.BORDER_CONSTANT, value=(1)) x1 -= pl y1 -= pt x2 += pr y2 += pb part_info['xyxy'] = [x1, y1, x2, y2] part_info['depth'] = depth part_info['img'] = img part_info['mask'] = (img[..., -1] > 10).astype(np.uint8) * 255 def load_parts(srcp, rotate=False, pad=0, min_width=64): srcimg = osp.join(srcp, 'src_img.png') fullpage = np.array(Image.open(srcimg).convert('RGBA')) infop = osp.join(srcp, 'info.json') infos = json2dict(infop) part_dict_list = [] tag2pd = {} part_id = 0 min_sz = 12 if isinstance(pad, int): pad = [pad] * 4 if rotate: fullpage = np.rot90(fullpage, 3, ) for tag, partdict in infos['parts'].items(): # img = Image.open(osp.join(srcp, tag + '.png')).convert('RGBA') # depthp = osp.join(srcp, tag + '_depth.png') # img = np.array(img) # p_test = max(img.shape[:2]) // 10 # mask = img[..., -1] > 10 # if np.sum(mask[:-p_test, :-p_test]) > 4: # if rotate: # img = np.rot90(img, 3) # mask = np.rot90(mask, 3, ) # xyxy = cv2.boundingRect(cv2.findNonZero(mask.astype(np.uint8))) # xyxy = np.array(xyxy) # h, w = xyxy[2:] # xyxy[2] += xyxy[0] # xyxy[3] += xyxy[1] # p = min_width - w # if p > 0: # if xyxy[0] >= p: # xyxy[0] -= p # else: # xyxy[2] += p # p = min_sz - h # if p > 0: # if xyxy[1] >= p: # xyxy[1] -= p # else: # xyxy[3] += p # depth = np.array(Image.open(depthp).convert('L')) # if rotate: # depth = np.rot90(depth, 3) # x1, y1, x2, y2 = xyxy # mask = mask[y1: y2, x1: x2].copy() # img = img[y1: y2, x1: x2].copy() # depth = depth[y1: y2, x1: x2].copy() # pt, pb, pl, pr = pad # if pt > 0 or pb > 0 or pl > 0 or pr > 0: # img = cv2.copyMakeBorder(img, pt, pb, pl, pr, cv2.BORDER_CONSTANT, value=(0, 0, 0, 0)) # depth = cv2.copyMakeBorder(depth, pt, pb, pl, pr, cv2.BORDER_CONSTANT, value=(255)) # mask = cv2.copyMakeBorder(mask.astype(np.uint8), pt, pb, pl, pr, cv2.BORDER_CONSTANT, value=(0)) > 0 # x1 -= pl # y1 -= pt # x2 += pr # y2 += pb # xyxy = [x1, y1, x2, y2] # dmin, dmax = partdict['depth_min'], partdict['depth_max'] # depth = np.array(depth, dtype=np.float32) / 255 * (dmax - dmin) + dmin p = load_part(osp.join(srcp, tag + '.png'), rotate=rotate, pad=pad, min_width=min_width, min_sz=min_sz) if p is not None: tag2pd[tag] = p tag2pd[tag]['part_id'] = part_id part_dict_list.append(tag2pd[tag]) part_id += 1 return fullpage, infos, part_dict_list