import argparse
import collections
import datetime
import glob
import json
import os
import os.path as osp
import sys
import uuid
import imgviz
import cv2
import numpy as np
import random
import shutil
from tqdm import tqdm
import labelme
import pycocotools.mask
def parse_opt():
parser = argparse.ArgumentParser()
parser.add_argument("--img_dir", default = r"图片以及Json文档保存位置的上级目录", help="input annotated directory")
parser.add_argument("--Folders", default=["图片以及Json文档保存文件名"])
parser.add_argument("--ClsLabels", default = [''], help="labels file") # '标签'
parser.add_argument("--YOLOBoxOnly", default=True, help = " yolov8需要Segmetation一起训练 ")
parser.add_argument("--Cencer_Points_Format", default = True, action="store_true")
parser.add_argument("--threshold_and", type=int, default=18)
parser.add_argument("--threshold_or", type=int, default=7)
parser.add_argument("--show_cencer_points", default = False, action="store_true")
parser.add_argument("--EmptyPointFolder", default=r"D:\Desktop\EmptyPointFolder", action="store_true")
args = parser.parse_args()
return args
def main(args):
for Folder in args.Folders:
LabelNum = len(args.ClsLabels)
output_dir = os.path.join(args.img_dir, Folder, "labels")
ImgPath = os.path.join(args.img_dir, Folder, "images")
jsonPath = ImgPath
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if not os.path.exists(args.EmptyPointFolder):
os.makedirs(args.EmptyPointFolder)
class_name_to_id = {}
if args.ClsLabels is str and args.ClsLabels.endwiths(".txt"):
for i, line in enumerate(open(args.ClsLabels).readlines()):
class_id = i
class_name = line.strip()
if class_id == -1:
assert class_name == "__ignore__"
continue
class_name_to_id[class_name] = class_id
else:
for id, name in enumerate(args.ClsLabels):
class_name_to_id[name] = id
ImgItems = glob.glob(osp.join(ImgPath , "*.jpg"))+glob.glob(osp.join(ImgPath , "*.bmp"))+glob.glob(osp.join(ImgPath , "*.png"))
for imgPath in tqdm(ImgItems):
img = cv2.imread(imgPath)
h, w = img.shape[0], img.shape[1]
JsonPath = imgPath.split(".")[0]+".json"
if not os.path.exists(JsonPath):
out_txt_path = os.path.join(output_dir, os.path.basename(imgPath).split(".")[0]+".txt")
out_txt_f = open(out_txt_path, 'w')
out_txt_f.write('\n')
continue
data = json.load(open(JsonPath))
# print(JsonPath)
base = os.path.splitext(os.path.basename(JsonPath))[0]
polygonData ={}
EmptyPoint = 0
box_list = []
seg_list = []
point_list = []
# height = data['imageHeight']
# width = data['imageWidth']
for id, shape in enumerate(data['shapes']):
label_name = shape['label']
if shape['shape_type'] == 'rectangle':
box_list.append(shape)
elif shape['shape_type'] == 'polygon':
if label_name not in class_name_to_id:
shape['label'] = label_name
print(label_name, "not in ClsLabels")
continue
seg_list.append(shape)
elif shape['shape_type'] == 'point':
point_list.append(shape)
if len(box_list):
with open(os.path.join(output_dir, (base + ".txt")), 'w') as file:
for box_shape, point_shape in zip(box_list,point_list):
label_name = box_shape['label']
if label_name not in class_name_to_id:
print(label_name, "not in ClsLabels")
continue
cls_id = class_name_to_id[label_name]
label_point = box_shape['points']
label_point = np.array(label_point).astype(np.int32)
# print(label_point)
X0,Y0,X1,Y1 = label_point[0][0], label_point[0][1], label_point[1][0], label_point[1][1]
Xmin , Ymin = min(X0, X1), min(Y0, Y1)
Xmax , Ymax = max(X0, X1), max(Y0, Y1)
Xmin,Ymin,Xmax,Ymax = max(int(Xmin),0), max(int(Ymin),0), min(max(int(Xmax),0),w) , min(max(int(Ymax),0),h)
if (Xmax-Xmin)<args.threshold_and and (Ymax-Ymin)<args.threshold_and:
continue
if (Xmax-Xmin)<args.threshold_or or (Ymax-Ymin)<args.threshold_or:
continue
b_w = (Xmax - Xmin)
b_h = (Ymax - Ymin)
b_x = Xmin + b_w / 2
b_y = Ymin + b_h / 2
b = [cls_id] + [b_x / w, b_y / h, b_w / w, b_h / h]
KeyPoint = np.array(point_shape['points']).flatten().astype(np.int32)
KeyPoint_X, KeyPoint_Y = KeyPoint[0], KeyPoint[1]
c_x = round((KeyPoint_X / w), 3)
c_y = round((KeyPoint_Y / h), 3)
if args.Cencer_Points_Format:
box = b + [c_x] + [c_y]
else:
box = b
line = *(box), # cls, box or segments
file.write(('%g ' * len(line)).rstrip() % line + '\n')
if args.show_cencer_points:
cv2.circle(img, (KeyPoint_X, KeyPoint_Y), radius=2, color=(0, 0, 255), thickness=2,
lineType=cv2.LINE_AA)
file.close()
if args.show_cencer_points:
cv2.imwrite(os.path.join(output_dir, (base + ".png")), img)
elif len(seg_list):
for id, seg_shape in enumerate(seg_list):
ann = {}
label_name = seg_shape['label']
cls_id = class_name_to_id[label_name]
points = seg_shape["points"]
points = np.asarray(points).reshape((-1, 1, 2))
if args.YOLOBoxOnly:
coordinates = np.array(points).reshape(-1, 2)
x0 = np.min(coordinates[:, 0])
x1 = np.max(coordinates[:, 0])
y0 = np.min(coordinates[:, 1])
y1 = np.max(coordinates[:, 1])
b_w = (x1 - x0)
b_h = (y1 - y0)
b_x = x0 + b_w / 2
b_y = y0 + b_h / 2
b = [cls_id] + [b_x / w, b_y / h, b_w / w, b_h / h]
ann["box"] = b
ann["seg_int32"] = np.array(points).reshape(-1, 2)
s = (np.array(points).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
if not args.Cencer_Points_Format:
ann["segment"] = [cls_id] + s
else:
ann["segment"] = s
ann["cls_id"] = [cls_id]
polygonData[id] = ann
if args.Cencer_Points_Format:
for point_shape in point_list:
points = point_shape["points"]
points = np.asarray(points).flatten().tolist()
for id, ann in polygonData.items():
seg = ann['seg_int32'].astype(np.int32)
result = cv2.pointPolygonTest(seg, tuple(points), False)
if result >= 0:
ann["center"]=points
c_x = round((ann["center"][0] / w), 3)
c_y = round((ann["center"][1] / h), 3)
ann["Cencer_Points"] = [c_x] + [c_y]
if args.show_cencer_points:
cv2.circle(img, (int(ann["center"][0]),int(ann["center"][1])), radius=2, color=(0,0,255),thickness=2,lineType=cv2.LINE_AA)
for id, ann in polygonData.items():
if args.Cencer_Points_Format:
if "Cencer_Points" not in ann.keys():
shutil.move(imgPath, os.path.join(args.EmptyPointFolder, os.path.basename(imgPath)))
shutil.move(JsonPath, os.path.join(args.EmptyPointFolder, os.path.basename(JsonPath)))
EmptyPoint = 1
break
if not EmptyPoint:
with open(os.path.join(output_dir, (base + ".txt")), 'w') as file:
for id, ann in polygonData.items():
if args.YOLOBoxOnly:
if args.Cencer_Points_Format:
box = ann["box"]+ann["Cencer_Points"]
else:
box = ann["box"]
line = *(box), # cls, box or segments
file.write(('%g ' * len(line)).rstrip() % line+'\n')
else:
if args.Cencer_Points_Format:
segment = ann["Cencer_Points"] + ann["cls_id"] + ann["segment"]
else:
segment = ann["segment"]
line = *(segment), # cls, box or segments
file.write(('%g ' * len(line)).rstrip() % line + '\n')
file.close()
if args.show_cencer_points:
cv2.imwrite(os.path.join(output_dir,(base+".png")),img)
else:
with open(os.path.join(output_dir, (base + ".txt")), 'w') as file:
pass
file.close()
EmptyFiles = os.listdir(args.EmptyPointFolder)
if len(EmptyFiles)==0:
os.rmdir(args.EmptyPointFolder)
if __name__ == "__main__":
args = parse_opt()
main(args)
2228
评论
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
查看更多评论
添加红包
