-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathprepare_Terraref_testset.py
104 lines (76 loc) · 3.81 KB
/
prepare_Terraref_testset.py
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
import os
import shutil
import cv2
import numpy as np
from tqdm import tqdm
from utils.utils import read_label, makedir, get_layout_image
#Path of Global Wheat Head Detection Dataset
gwhd_2021_path = 'datasets/gwhd_2021/'
#Path to save the images of Terraref domain
Terraref_path = 'datasets/gwhd_2021/Terraref/target/'
makedir(Terraref_path)
Terraref_ori_source_path = 'datasets/gwhd_2021/Terraref/source/'
makedir(Terraref_ori_source_path)
Terraref_cut_Path = 'datasets/gwhd_2021/Terraref_x9/'
makedir(Terraref_cut_Path)
Terraref_source_Path = Terraref_cut_Path + 'source/'
makedir(Terraref_source_Path)
Terraref_target_Path = Terraref_cut_Path + 'target/'
makedir(Terraref_target_Path)
gwhd_2021_img_path = gwhd_2021_path + 'images/'
csvFiles = gwhd_2021_path + 'competition_test.csv'
Terraref_label = open(Terraref_cut_Path + 'label_x9.txt', 'w', encoding='utf-8')
names_img_with_wheat = open(Terraref_cut_Path + 'with_wheat.txt', 'w', encoding='utf-8')
names_img_wo_wheat = open(Terraref_cut_Path + 'wo_wheat.txt', 'w', encoding='utf-8')
#read labels of testset
label_dic = read_label(csvFiles)
#devide images belonging to Terraref
for img_name in os.listdir(gwhd_2021_img_path):
if img_name in label_dic.keys():
domain = label_dic[img_name][-1]
if 'Terraref' in domain:
shutil.copy2(gwhd_2021_img_path + img_name, Terraref_path)
for imgName in tqdm(os.listdir(Terraref_path)):
if '.png' not in imgName:
continue
img = cv2.imread(Terraref_path + imgName)
if imgName in label_dic.keys() :
laBel = label_dic[imgName][0]
imgName = imgName[:-4]
laBel = laBel.replace(" ", ",").split(';')
source_img, n_lays = get_layout_image(img, laBel)
cv2.imwrite(Terraref_ori_source_path + imgName + '.png', source_img)
if 'no_box' not in laBel:
orginal_bbox = np.array([np.array(list(map(int,bbox.split(',')))) for bbox in laBel])
#cut image into 512*512, with step=256(x9), and get layout images and new labels(txt)
i = 0
for y_i in range(3):
for x_i in range(3):
x_start = x_i*256
y_start = y_i*256
img_cut = img[y_start:y_start+512, x_start:x_start+512,:]
source_img_cut = source_img[y_start:y_start+512, x_start:x_start+512,:]
cv2.imwrite(Terraref_target_Path + imgName + '_' + str(i) + '.png', img_cut)
cv2.imwrite(Terraref_source_Path + imgName + '_' + str(i) + '.png', source_img_cut)
Terraref_label.write('%s.png'%(Terraref_target_Path + imgName + '_' + str(i)))
if 'no_box' not in laBel:
bbox=orginal_bbox.copy()
bbox[:, 0] = bbox[:, 0] - x_start
bbox[:, 2] = bbox[:, 2] - x_start
bbox[:, 1] = bbox[:, 1] - y_start
bbox[:, 3] = bbox[:, 3] - y_start
bbox[:, 0:2][bbox[:, 0:2]<0] = 0
bbox[:, 2][bbox[:, 2]>511] = 511
bbox[:, 3][bbox[:, 3]>511] = 511
bboxes = bbox[np.logical_and((bbox[:, 2] - bbox[:, 0])>10, (bbox[:, 3] - bbox[:, 1])>10)] # discard invalid box
if bboxes.shape[0] > 0:
names_img_with_wheat.write('%s.png'%(imgName + '_' + str(i)) + '\n')
for bbox in bboxes:
bbox = ','.join(map(str, bbox))
Terraref_label.write(' ' + bbox + ',0')
else:
names_img_wo_wheat.write('%s.png'%(imgName + '_' + str(i)) + '\n')
else:
names_img_wo_wheat.write('%s.png'%(imgName + '_' + str(i)) + '\n')
Terraref_label.write('\n')
i += 1