wehub-resource-sync
105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
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)
|