table_structure_recognization/preprocess.py at master · kapitsa2811/table_structure_recognization · GitHub

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
from utils import preprocess, write_filename, split_train_val
import os
import cv2


dir_path = r"D:\tablebank\ICDAR2019_cTDaR\ICDAR2019_cTDaR\training\TRACKB1\ground_truth"
save_path = r"D:\tablebank\ICDAR2019_cTDaR\ICDAR2019_cTDaR\training\TRACKB1\preprocess"
files =  os.listdir(dir_path)

image_name = []
for img_file in files:
    name = os.path.splitext(img_file)
    print("current file %s" %(name[0]))

    if name[1] == '.jpg' or name[1] == '.png':
        img_path = os.path.join(dir_path, img_file)
        xml_path = os.path.join(dir_path, name[0]+'.xml')
        img_table, fm_table = preprocess(img_path, xml_path)
        assert len(img_table)==len(fm_table), "image and feature map not match"

        for idx in range(len(img_table)):
            image_name.append(name[0]+ '_'+ str(idx))
            cv2.imwrite(os.path.join(save_path, name[0]+ '_'+ str(idx)+ '.jpg'), img_table[idx])
            cv2.imwrite(os.path.join(save_path, name[0]+ '_'+ str(idx)+ '_maskh.png'), fm_table[idx][0])
            cv2.imwrite(os.path.join(save_path, name[0]+ '_'+ str(idx)+ '_maskv.png'), fm_table[idx][1])
    if name[1] == '.tiff':
        img_path = os.path.join(dir_path, img_file)
        tif = TIFF.open(img_path, mode='r')
        for img in list(tif.iter_images()):
            tiff = img
        cv2.imwrite(os.path.join(dir_path, name[0]+'.jpg'))

        img_path = os.path.join(dir_path, name[0]+'.jpg')
        xml_path = os.path.join(dir_path, name[0]+'.xml')
        img_table, fm_table = preprocess(img_path, xml_path)
        assert len(img_table)==len(fm_table), "image and feature map not match"

        for idx in range(len(img_table)):
            image_name.append(name[0]+ '_'+ str(idx))
            cv2.imwrite(os.path.join(save_path, name[0]+ '_'+ str(idx)+ '.jpg'), img_table[idx])
            cv2.imwrite(os.path.join(save_path, name[0]+ '_'+ str(idx)+ '_maskh.png'), fm_table[idx][0])
            cv2.imwrite(os.path.join(save_path, name[0]+ '_'+ str(idx)+ '_maskv.png'), fm_table[idx][1])

dir_path = r'D:\a上财学习材料\20毕业论文\实验\data'
filename_path = 'cTDaR19_imagename.txt'
write_filename(os.path.join(dir_path, filename_path), image_name)

# split train and test
split_train_val(dir_path, filename_path)