Detectron：训练自己的数据集——将自己的数据格式转换成COCO格式

# coding=utf-8 import json import os import cv2 import time start = time.time() # 根路径，里面包含images(图片文件夹)，annos.txt(bbox标注)，classes.txt(类别标签),以及annotations文件夹(如果没有则会自动创建，用于保存最后的json) root_path = '/home/renjiaxin/detectron/detectron/datasets/data/coco/' # 用于创建训练集或验证集 phase = 'train' # 训练集和验证集划分的界线 # split = 7955 dataset = {'licenses': [], 'info': {}, 'categories': [], 'images': [], 'annotations': []} # 打开类别标签 with open(os.path.join(root_path, 'classes.txt')) as f: classes = f.read().strip().split() # 建立类别标签和数字id的对应关系 for i, cls in enumerate(classes, 1): dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'mark'}) # 读取images文件夹的图片名称 _indexes = [f for f in os.listdir(os.path.join(root_path, 'aug'))] # # 判断是建立训练集还是验证集 # if phase == 'train': # indexes = [line for i, line in enumerate(_indexes, 1) if i <= split] # elif phase == 'val': # indexes = [line for i, line in enumerate(_indexes, 1) if i > split] # 读取Bbox信息 with open(os.path.join(root_path, 'total_train_annos.txt')) as tr: annos = tr.readlines() n_p , n_a = 0, 0 for k, index in enumerate(_indexes): # 用opencv读取图片，得到图像的宽和高 im = cv2.imread(os.path.join(root_path, 'aug/') + index) height, width, _ = im.shape if phase == 'train': n_p += 1 print '第{}张图片'.format(n_p) # 添加图像的信息到dataset中 dataset['images'].append({'file_name': index, 'id': k, 'width': width, 'height': height}) elif phase == 'val': # 添加图像的信息到dataset中 dataset['images'].append({'file_name': index, 'id': k + split + 10000000, # 加上一个很大的数是为了和增强后的训练集进行区分， # 不然会有重叠，下面annotations的image_id对应着这个ID，所以也是这个原因这两者的数值相同， # 如果不增强训练集的话，这个常数改为1，并且这个常数只在val中需要添加，train不用管 'width': width, 'height': height}) for ii, anno in enumerate(annos): parts = anno.strip().split() # 如果图像的名称和标记的名称对上，则添加标记 if parts[0] == index: n_a += 1 print '第{}个注释'.format(n_a) # 类别 cls_id = parts[1] # x_min x1 = float(parts[2]) # y_min y1 = float(parts[3]) # x_max x2 = float(parts[4]) # y_max y2 = float(parts[5]) width = max(0, x2 - x1) height = max(0, y2 - y1) if phase == 'train': dataset['annotations'].append({ 'area': width * height, 'bbox': [x1, y1, width, height], 'category_id': int(cls_id), 'id': ii, 'image_id': k, 'iscrowd': 0, # mask, 矩形是从左上角点按顺时针的四个顶点 'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]] }) elif phase == 'val': dataset['annotations'].append({ 'area': width * height, 'bbox': [x1, y1, width, height], 'category_id': int(cls_id), 'id': ii + 10000000, # 默认为1 'image_id': split + k + 10000000, # 该常数默认为1 'iscrowd': 0, # mask, 矩形是从左上角点按顺时针的四个顶点 'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]] }) # 保存结果的文件夹 folder = os.path.join(root_path, 'annotations') if not os.path.exists(folder): os.makedirs(folder) json_name = os.path.join(root_path, 'annotations/2{}.json'.format(phase)) with open(json_name, 'w') as f: json.dump(dataset, f) end = time.time() print "time:{}小时".format((end-start)/3600)