Spaces:
Sleeping
Sleeping
File size: 4,208 Bytes
6fc683c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
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)
|