122 lines
4.7 KiB
Python
122 lines
4.7 KiB
Python
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]
|