-
Notifications
You must be signed in to change notification settings - Fork 237
/
Copy pathjson2yolo.py
121 lines (112 loc) · 4.85 KB
/
json2yolo.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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import json
import os
import numpy as np
from copy import deepcopy
import cv2
def allFilePath(rootPath,allFIleList):
fileList = os.listdir(rootPath)
for temp in fileList:
if os.path.isfile(os.path.join(rootPath,temp)):
allFIleList.append(os.path.join(rootPath,temp))
else:
allFilePath(os.path.join(rootPath,temp),allFIleList)
def xywh2yolo(rect,landmarks_sort,img):
h,w,c =img.shape
rect[0] = max(0, rect[0])
rect[1] = max(0, rect[1])
rect[2] = min(w - 1, rect[2]-rect[0])
rect[3] = min(h - 1, rect[3]-rect[1])
annotation = np.zeros((1, 12))
annotation[0, 0] = (rect[0] + rect[2] / 2) / w # cx
annotation[0, 1] = (rect[1] + rect[3] / 2) / h # cy
annotation[0, 2] = rect[2] / w # w
annotation[0, 3] = rect[3] / h # h
annotation[0, 4] = landmarks_sort[0][0] / w # l0_x
annotation[0, 5] = landmarks_sort[0][1] / h # l0_y
annotation[0, 6] = landmarks_sort[1][0] / w # l1_x
annotation[0, 7] = landmarks_sort[1][1] / h # l1_y
annotation[0, 8] = landmarks_sort[2][0] / w # l2_x
annotation[0, 9] = landmarks_sort[2][1] / h # l2_y
annotation[0, 10] = landmarks_sort[3][0] / w # l3_x
annotation[0, 11] = landmarks_sort[3][1] / h # l3_y
# annotation[0, 12] = (landmarks_sort[0][0]+landmarks_sort[1][0])/2 / w # l4_x
# annotation[0, 13] = (landmarks_sort[0][1]+landmarks_sort[1][1])/2 / h # l4_y
return annotation
def order_points(pts):
rect = np.zeros((4, 2), dtype = "float32")
s = pts.sum(axis = 1)
rect[0] = pts[np.argmin(s)]
rect[2] = pts[np.argmax(s)]
diff = np.diff(pts, axis = 1)
rect[1] = pts[np.argmin(diff)]
rect[3] = pts[np.argmax(diff)]
# return the ordered coordinates
return rect
def four_point_transform(image, pts):
rect = order_points(pts)
(tl, tr, br, bl) = rect
widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
maxWidth = max(int(widthA), int(widthB))
heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
maxHeight = max(int(heightA), int(heightB))
dst = np.array([
[0, 0],
[maxWidth - 1, 0],
[maxWidth - 1, maxHeight - 1],
[0, maxHeight - 1]], dtype = "float32")
M = cv2.getPerspectiveTransform(rect, dst)
warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
# return the warped image
return warped
if __name__ == "__main__":
pic_file_list = []
pic_file = r"/mnt/Gpan/Mydata/pytorchPorject/datasets/ccpd/train_bisai/train_bisai"
save_small_path = "small"
label_file = ['0','1']
allFilePath(pic_file,pic_file_list)
count=0
index = 0
for pic_ in pic_file_list:
if not pic_.endswith(".jpg"):
continue
count+=1
img = cv2.imread(pic_)
img_name = os.path.basename(pic_)
txt_name = img_name.replace(".jpg",".txt")
txt_path = os.path.join(pic_file,txt_name)
json_file_ = pic_.replace(".jpg",".json")
if not os.path.exists(json_file_):
continue
with open(json_file_, 'r',encoding='utf-8') as a:
data_dict = json.load(a)
# print(data_dict['shapes'])
with open(txt_path,"w") as f:
for data_message in data_dict['shapes']:
index+=1
label=data_message['label']
points = data_message['points']
pts = np.array(points)
# pts=order_points(pts)
# new_img = four_point_transform(img,pts)
roi_img_name = label+"_"+str(index)+".jpg"
save_path=os.path.join(save_small_path,roi_img_name)
# cv2.imwrite(save_path,new_img)
x_max,y_max = np.max(pts,axis=0)
x_min,y_min = np.min(pts,axis=0)
rect = [x_min,y_min,x_max,y_max]
rect1=deepcopy(rect)
annotation=xywh2yolo(rect1,pts,img)
print(data_message)
label = data_message['label']
str_label = label_file.index(label)
# str_label = "0 "
str_label = str(str_label)+" "
for i in range(len(annotation[0])):
str_label = str_label + " " + str(annotation[0][i])
str_label = str_label.replace('[', '').replace(']', '')
str_label = str_label.replace(',', '') + '\n'
f.write(str_label)
print(count,img_name)
# point=data_message[points]