import os import json import glob import random import cv2 def read_txt(file): with open(file, "r", encoding="utf8") as f: data = [line.strip() for line in f] return data def write_txt(file, data): with open(file, "w", encoding="utf8") as f: for item in data: f.write(item + "\n") def write_json(file, data): with open(file, "w", encoding="utf8") as f: json.dump(data, f, ensure_ascii=False, sort_keys=True) def read_json(file): with open(file, "r", encoding="utf8") as f: data = json.load(f) return data def get_colors(kie_labels): random.seed(1997) colors = [] for _ in range(len(kie_labels)): color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) colors.append(color) return colors def normalize_box(box, width, height): assert ( max(box) <= width or max(box) <= height ), "box must smaller than width, height; max box = {}, width = {}, height = {}".format( max(box), width, height ) return [ int(1000 * (box[0] / width)), int(1000 * (box[1] / height)), int(1000 * (box[2] / width)), int(1000 * (box[3] / height)), ] def unnormalize_box(bbox, width, height): return [ width * (bbox[0] / 1000), height * (bbox[1] / 1000), width * (bbox[2] / 1000), height * (bbox[3] / 1000), ] def load_image_paths_and_labels(data_dir): r"""Load (image path, label) pairs into a DataFrame with keys ``image_path`` and ``label`` @todo Add OCR paths here """ img_paths = [path for path in glob.glob(data_dir + "/*") if ".txt" not in path] label_paths = [os.path.splitext(path)[0] + ".txt" for path in img_paths] return img_paths, label_paths import cv2 def read_image_file(img_path): image = cv2.imread(img_path) return image def normalize_bbox(x1, y1, x2, y2, w, h): x1 = int(float(min(max(0, x1), w))) x2 = int(float(min(max(0, x2), w))) y1 = int(float(min(max(0, y1), h))) y2 = int(float(min(max(0, y2), h))) return (x1, y1, x2, y2) def extend_crop_img( left, top, right, bottom, margin_l=0, margin_t=0.03, margin_r=0.02, margin_b=0.05 ): top = top - (bottom - top) * margin_t bottom = bottom + (bottom - top) * margin_b left = left - (right - left) * margin_l right = right + (right - left) * margin_r return left, top, right, bottom def get_crop_img_and_bbox(img, bbox, extend: bool): """ img : numpy array img bbox : should be xyxy format """ if len(bbox) == 5: left, top, right, bottom, _conf = bbox elif len(bbox) == 4: left, top, right, bottom = bbox if extend: left, top, right, bottom = extend_crop_img(left, top, right, bottom) left, top, right, bottom = normalize_bbox( left, top, right, bottom, img.shape[1], img.shape[0] ) assert (bottom - top) * (right - left) > 0, "bbox is invalid" crop_img = img[top:bottom, left:right] return crop_img, (left, top, right, bottom) import json import os def load_kie_labels_yolo(label_path): with open(label_path, 'r') as f: lines = f.read().splitlines() words, boxes, labels = [], [], [] for line in lines: x1, y1, x2, y2, text, kie = line.split("\t") x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) if text != " ": words.append(text) boxes.append((x1, y1, x2, y2)) labels.append(kie) return words, boxes, labels def create_empty_kie_dict(): return {cfg.KIE_LABELS[i]: [] for i in range(len(cfg.KIE_LABELS)) if cfg.KIE_LABELS[i] != cfg.IGNORE_KIE_LABEL} def write_to_json_(file_path, content): with open(file_path, mode='w', encoding='utf8') as f: json.dump(content, f, ensure_ascii=False) def read_json(file_path): with open(file_path, 'r') as f: return json.load(f) def get_name(file_path, ext: bool = True): file_path_ = os.path.basename(file_path) return file_path_ if ext else os.path.splitext(file_path_)[0] def construct_file_path(dir, file_path, ext=''): ''' args: dir: /path/to/dir file_path /example_path/to/file.txt ext = '.json' return /path/to/dir/file.json ''' return os.path.join( dir, get_name(file_path, True)) if ext == '' else os.path.join( dir, get_name(file_path, False)) + ext def write_to_txt_(file_path, content): with open(file_path, 'w') as f: f.write(content)