diff --git a/label_anything/tools/convert_to_mask_polygon_coco.py b/label_anything/tools/convert_to_mask_polygon_coco.py new file mode 100644 index 00000000..a556402a --- /dev/null +++ b/label_anything/tools/convert_to_mask_polygon_coco.py @@ -0,0 +1,218 @@ +import json +import os +import sys +import cv2 +import argparse +import numpy as np +from tqdm import tqdm +from itertools import groupby +from label_studio_converter.brush import decode_rle + +import jinja2 +from ../sam/filter_poly import NearNeighborRemover +import shutil + +def parse_args(): + parser = argparse.ArgumentParser(description='Label studio convert to Coco fomat') + parser.add_argument('--json_file_path',default='project.json', help='label studio output json') + parser.add_argument('--out_dir',default='../mmdetection/data/my_set', help='output dir of Coco format json') + parser.add_argument('--classes',default=None, help='Classes list of the dataset, if None please check the output.') + parser.add_argument('--out_config',default=None, choices=['mask-rcnn_r50_fpn','rtmdet-ins_s',None],help='config mode') + + args = parser.parse_args() + return args + + +coco_format = {} +images = [] +categories = [] +annotations = [] +bbox = [] +area = [] + +def binary_mask_to_rle(binary_mask): + rle = {'counts': [], 'size': list(binary_mask.shape)} + counts = rle.get('counts') + for i, (value, elements) in enumerate(groupby(binary_mask.ravel(order='F'))): + if i == 0 and value == 1: + counts.append(0) + counts.append(len(list(elements))) + return rle + +def rle2mask(rle,height, width): + """Convert RLE to mask + + :param rle: list of ints in RLE format + :param height: height of binary mask + :param width: width of binary mask + :return: binary mask + """ + mask = decode_rle(rle) + mask = mask.reshape(-1,4)[:, 0] + if mask.max()==255: + mask=np.where(mask==255,1,0) + + return mask.reshape((height, width)) + +def format_to_coco(args): + # 读取label studio格式的JSON文件 + json_file_path=args.json_file_path + + with open(json_file_path, 'r') as file: + contents = json.loads(file.read()) + if sys.platform == 'linux': + image_path_from=os.path.join('~/.local/share/label-studio/media/upload/',os.path.dirname(contents[0]['data']['image']).split('/')[-1]) + elif sys.platform == 'win32': + print("JSON Contents:", contents[0]) + print("JSON Contents:", contents[0]['data']['image']) + print("JSON Contents:", contents[0]['data']['image'].split('/')[-1]) + + #image_path_from=os.path.join('~/AppData/Local/label-studio/label-studio/media/upload/',os.path.dirname(contents[0]['data']['image']).split('/')[-1]) + image_path_from=os.path.join('~/AppData/Local/label-studio/label-studio/media/upload/',os.path.dirname(contents[0]['data']['image']).split('/')[-1]) + else: + raise ValueError("The system does not support.") + image_path_to=args.out_dir + # 将coco格式保存到文件中 + output_dir=image_path_to + output_ann_path=os.path.join(output_dir,'annotations') + os.makedirs(output_ann_path, exist_ok=True) + os.makedirs(os.path.join(image_path_to,'images'), exist_ok=True) + # 初始化coco格式的字典 + coco_format = { + "images": [], + "categories": [], + "annotations": [] + } + + #自定义类别标签顺序 + if args.classes is not None: + for category in args.classes: + coco_format["categories"].append({ + "id": len(coco_format["categories"]), + "name": category + }) + + index_cnt=0 + # 遍历每个标注 + for index_annotation in tqdm(range(len(contents))): + result = contents[index_annotation]['annotations'][0]['was_cancelled'] + + if result == False: + width_from_json = contents[index_annotation]["annotations"][0]["result"][0]["original_width"] + height_from_json = contents[index_annotation]["annotations"][0]["result"][0]["original_height"] + + labels = contents[index_annotation]['annotations'][0]['result'] + + for i in range(len(labels)): + + if 'rle' in labels[i]['value'].keys(): + label_rel=labels[i] + else: + continue + + image_id = contents[index_annotation]["id"] + image_json_name_ = contents[index_annotation]["data"]['image'].split('/')[-1] + image_json_name = image_json_name_.split('-', 1)[1] + try: + category = label_rel['value']['brushlabels'][0] + except: + category = label_rel['value']['rectanglelabels'][0] + rle = label_rel['value']['rle'] + mask=rle2mask(rle,height_from_json,width_from_json) #由rle转换为mask + contours, hierarchy = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) #处理mask获取polygon + point_list = [] + for contour in contours: + points = [] + for point in contour: + x, y = point[0] + points.append([float(x), float(y)]) + point_list.extend(points) + filter_points=NearNeighborRemover(distance_threshold=0.4).remove_near_neighbors(point_list) + polygon_points=[] + polygon_points[0]=[] + for point in filter_points: #coco数据集polygon格式 + polygon_points[0].append(point) + new_contours = [] + for contour in contours: + new_contours.extend(list(contour)) + new_contours = np.array(new_contours) + x, y, w, h = cv2.boundingRect(new_contours) + + bbox = [x, y, w, h] + area = w * h + + # 添加图像信息到coco格式 + coco_format["images"].append({ + "id": image_id, + "file_name": image_json_name, + "width": width_from_json, + "height": height_from_json + }) + + # 检查类别是否已经添加到categories变量中,如果没有,将其添加到categories变量 + if not any(d["name"] == category for d in coco_format["categories"]): + coco_format["categories"].append({ + "id": len(coco_format["categories"]), + "name": category + }) + + # 添加标注信息到coco格式 + coco_format["annotations"].append({ + "id": index_cnt, + "image_id": image_id, + "category_id": [d["id"] for d in coco_format["categories"] if d["name"] == category][0], + "segmentation": polygon_points, + "bbox": bbox, + "area": area, + 'iscrowd': 0#polygon format + + }) + index_cnt+=1 + image_from=os.path.join(image_path_from,image_json_name_) + image_to=os.path.join(image_path_to,'images',image_json_name) + shutil.copy2(os.path.expanduser(image_from), image_to) + + classes_output=[d["name"] for d in coco_format["categories"]] + print(classes_output) + args.train_ann_file=os.path.join(output_ann_path,'ann.json') + with open(os.path.join(output_ann_path,'ann.json'), "w") as out_file: + json.dump(coco_format, out_file, ensure_ascii=False, indent=4) + return classes_output,args + +def move_to_cfg(args,classes_list): + if 'rtmdet-ins_s' in args.out_config: + config_path='config_template/rtmdet-ins_s_8xb32-300e_coco.py' + config_name='rtmdet-ins_s.py' + elif 'mask-rcnn_r50_fpn' in args.out_config: + config_path='config_template/mask-rcnn_r50_fpn_1x_coco.py' + config_name='mask-rcnn_r50_fpn.py' + + train_ann_pth='annotations/ann.json' + train_data_pf='images/' + num_classes = len(classes_list) + + data_root=str('\''+os.path.join('.',args.out_dir.split('/', 2)[-1]+'/')+'\'') + train_ann_file=val_ann_file=str('\''+train_ann_pth+'\'') + train_data_prefix=val_data_prefix=str('\''+train_data_pf+'\'') + variable_dict = {'class_name':tuple(classes_list), 'num_classes':num_classes, 'data_root':data_root,\ + 'train_ann_file':train_ann_file,'val_ann_file':train_ann_file, \ + 'train_data_prefix':train_data_prefix,'val_data_prefix':train_data_prefix} + current_dir = os.getcwd() + env = jinja2.Environment(loader=jinja2.FileSystemLoader(current_dir)) + temp = env.get_template(config_path) + + temp_out=temp.render(class_name=variable_dict['class_name'],num_classes=variable_dict['num_classes'],\ + data_root=variable_dict['data_root'],train_ann_file=variable_dict['train_ann_file'],\ + val_ann_file=variable_dict['val_ann_file'],train_data_prefix=variable_dict['train_data_prefix'],\ + val_data_prefix=variable_dict['val_data_prefix']) + with open(os.path.join(args.out_dir,config_name), "w", encoding='utf-8') as out_file: + out_file.writelines(temp_out) + out_file.close() + print(f'The config have been saved in \'{args.out_dir}\'') + +if __name__ == '__main__': + args = parse_args() + classes_output,args=format_to_coco(args) + if args.out_config is not None: + move_to_cfg(args,classes_output) + diff --git a/label_anything/tools/convert_to_rle_mask_coco.py b/label_anything/tools/convert_to_rle_mask_coco.py index 84cbf3e4..24ccc826 100644 --- a/label_anything/tools/convert_to_rle_mask_coco.py +++ b/label_anything/tools/convert_to_rle_mask_coco.py @@ -18,7 +18,7 @@ def parse_args(): parser.add_argument('--json_file_path',default='project.json', help='label studio output json') parser.add_argument('--out_dir',default='../mmdetection/data/my_set', help='output dir of Coco format json') parser.add_argument('--classes',default=None, help='Classes list of the dataset, if None please check the output.') - parser.add_argument('--out_config',default=None, choices=['mask-rcnn_r50_fpn','rtmdet-ins_s',None],help='config mode') + parser.add_argument('--out_config',default=mask-rcnn_r50_fpn, choices=['mask-rcnn_r50_fpn','rtmdet-ins_s',None],help='config mode') args = parser.parse_args() return args