Spaces:
Sleeping
Sleeping
| import os | |
| from PIL import Image | |
| import xml.etree.ElementTree as ET | |
| import numpy as np | |
| import json | |
| from PIL import Image | |
| from shutil import copyfile | |
| def convert(ROOT, TRACK, SPLIT): | |
| coco_data = { | |
| "images": [], | |
| "annotations": [], | |
| "categories": [{"id": 1, "name": "table"}, ], | |
| } | |
| DATA_DIR = f"{ROOT}/{TRACK}/{SPLIT}" | |
| prefix = "cTDaR_t0" if TRACK == "trackA_archival" else "cTDaR_t1" | |
| print(TRACK, SPLIT, prefix) | |
| table_count = 0 | |
| for file in sorted(os.listdir(DATA_DIR)): | |
| if file.startswith(prefix) and file.endswith(".jpg"): | |
| img = Image.open(os.path.join(DATA_DIR, file)) | |
| coco_data["images"].append( | |
| { | |
| "file_name": file, | |
| "height": img.height, | |
| "width": img.width, | |
| "id": int(file[7:-4]), | |
| } | |
| ) | |
| elif file.startswith(prefix) and file.endswith(".xml"): | |
| # print(file) | |
| tree = ET.parse(os.path.join(DATA_DIR, file)) | |
| root = tree.getroot() | |
| assert len(root.findall("./table/Coords")) > 0 | |
| for table_id in range(len(root.findall("./table/Coords"))): | |
| four_points = root.findall("./table/Coords")[table_id].attrib["points"] | |
| four_points = list(map(lambda x: x.split(","), four_points.split())) | |
| four_points = [[int(j) for j in i] for i in four_points] | |
| segmentation = [j for i in four_points for j in i] | |
| bbox = [ | |
| four_points[0][0], | |
| four_points[0][1], | |
| four_points[2][0] - four_points[0][0], | |
| four_points[2][1] - four_points[0][1], | |
| ] | |
| coco_data["annotations"].append( | |
| { | |
| "segmentation": [segmentation], | |
| "area": bbox[2] * bbox[3], | |
| "iscrowd": 0, | |
| "image_id": int(file[7:-4]), | |
| "bbox": bbox, | |
| "category_id": 1, | |
| "id": table_count, | |
| } | |
| ) | |
| table_count += 1 | |
| with open(f"{ROOT}/{TRACK}/{SPLIT}.json", "w") as f: | |
| json.dump(coco_data, f) | |
| def clean_img(DATA_DIR): | |
| for file in sorted(os.listdir(DATA_DIR)): | |
| if file.endswith(".JPG"): | |
| os.rename(os.path.join(DATA_DIR, file), os.path.join(DATA_DIR, file.replace(".JPG", ".jpg"))) | |
| elif file.endswith(".TIFF"): | |
| img = Image.open(os.path.join(DATA_DIR, file)) | |
| img.save(os.path.join(DATA_DIR, file.replace(".TIFF", ".jpg"))) | |
| os.remove(os.path.join(DATA_DIR, file)) | |
| elif file.endswith(".png"): | |
| img = Image.open(os.path.join(DATA_DIR, file)) | |
| img.save(os.path.join(DATA_DIR, file.replace(".png", ".jpg"))) | |
| os.remove(os.path.join(DATA_DIR, file)) | |
| if __name__ == '__main__': | |
| import argparse | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--root_dir', required=True) | |
| parser.add_argument('--target_dir', required=True) | |
| args = parser.parse_args() | |
| test_data_dir = os.path.join(args.root_dir, 'test', 'TRACKA') | |
| test_gt_dir = os.path.join(args.root_dir, 'test_ground_truth', 'TRACKA') | |
| training_data_dir = os.path.join(args.root_dir, 'training', 'TRACKA', 'ground_truth') | |
| raw_datas = {"train": [training_data_dir], "test": [test_data_dir, test_gt_dir]} | |
| TRACKS = ["trackA_modern", "trackA_archival"] | |
| SPLITS = ["train", "test"] | |
| for track in TRACKS: | |
| prefix = "cTDaR_t0" if track == "trackA_archival" else "cTDaR_t1" | |
| for split in SPLITS: | |
| os.makedirs(os.path.join(args.target_dir, track, split)) | |
| for source_dir in raw_datas[split]: | |
| for fn in os.listdir(source_dir): | |
| if fn.startswith(prefix): | |
| ffn = os.path.join(source_dir, fn) | |
| copyfile(ffn, os.path.join(args.target_dir, track, split, fn)) | |
| clean_img(os.path.join(args.target_dir, track, split)) | |
| convert(args.target_dir, track, split) | |