# 说明：
'''
1)military_object.names文件内容如下：
label1
label2
label3
2)此训练针对矩形框的标注
3）代码中很多键值都是自己造的，因为训练用不到这些(比如mask的训练)，仅仅只是为了跟coco格式保持一致
'''
import os
import json
import cv2
import random
import time
# coco格式最后储存位置
coco_format_save_path = '../datasets/military/annotations/instances_val_military.json'
# 类别文件，一行一个类
yolo_format_classes_path = '../datasets/military/annotations/military_object.names'
# yolo格式的注释文件
yolo_format_annotation_path = '../datasets/military/annotations/military_val_yolo.txt'
# 根据自己的数据集写类别。举个例子:
# categories_dict = [{'supercategory': 'None', 'id': 1, 'name': 'w3'},{'supercategory': 'None', 'id': 2, 'name': 'h3'}]
# 我有类别文件,本着能用代码绝不手写的原则
with open(yolo_format_classes_path,'r') as f1:
lines1 = f1.readlines()
categories = []
for j,label in enumerate(lines1):
label = label.strip()
categories.append({'id':j+1,'name':label,'supercategory': 'None'})
"""========================================后文就正式开始了========================================="""
write_json_context = dict()
write_json_context['info'] = {'description': '', 'url': '', 'version': '', 'year': 2020, 'contributor': '', 'date_created': '2020-06-16 11:00:08.5'}
write_json_context['licenses'] = [{'id': 1, 'name': None, 'url': None}]
write_json_context['categories'] = categories
write_json_context['images'] = []
write_json_context['annotations'] = []
# 前面都造好了基础(为了和coco一致，其实很多都用不到的)，现在开始解析自己的数据
with open(yolo_format_annotation_path,'r') as f2:
lines2 = f2.readlines()
# 第n个标注框计数,也就是annotations中的id
num_bboxes = 0
# 每一行就是一张图片的标注信息。
# 其他格式也很好做，这个遍历就是图片文件路径，后面bboxes的遍历就是解析当前图片的标注信息
# 如果xml或者json的,我的csdn中也有labelme标注的xml和json的解析代码(很久很久以前的代码,其实有更简洁的解析,但是我懒得写,哈哈)
# 看懂我的代码,改起来不要太简单
for i,line in enumerate(lines2):
img_context = {}
# 我的数据以空格分隔的,具体查看上面的截图
line = line.split(' ')
# 我在想:如果这张图没有任何标注,要不要保留呢,我选择保留负样本，当然你也可以打开下面这段代码舍弃
# if len(line) < 2:
# continue
# 文件名最好不要有空格，这是一种习惯
img_path = line[0].rstrip()
img_name = os.path.basename(img_path)
# 因为需要width和height,而我得yolo文件里没有，所以我还得读图片，很烦
# 我就用opencv读取了,当然用其他库也可以
# 别把图片路径搞错了，绝对路径和相对路径分清楚！
height,width = cv2.imread(img_path).shape[:2]
img_context['file_name'] = img_name
img_context['height'] = height
img_context['width'] = width
img_context['date_captured'] = '2020-06-16 11:00:08.5'
# 这么多id搞得我头都懵了,我猜这是第几张图序号吧,每行一张图，那当然就是第i张了
img_context['id'] = i
img_context['license'] = 1
img_context['coco_url'] =''
img_context['flickr_url'] = ''
write_json_context['images'].append(img_context)
# 这个地方如果有标注框继续，没有的话跳过，跟上面的区别在于是否将图片加入到images中
# 如果images中有这张图但是没有标注信息，那就是负样本，反之亦然
if len(line) < 2:
continue
for bbox in line[1:]:
bbox_dict = {}
xmin,ymin,xmax,ymax,class_id = bbox.strip().split(',')
# 我就有时候int和str不注意各种报错
xmin,ymin,xmax,ymax,class_id= float(xmin),float(ymin),float(xmax),float(ymax),int(class_id)
bbox_dict['id'] = num_bboxes
bbox_dict['image_id'] = i
bbox_dict['category_id'] = class_id + 1
bbox_dict['iscrowd'] = 0 # 前面有解释
h,w = abs(ymax-ymin),abs(xmax-xmin)
bbox_dict['area'] = h * w
bbox_dict['bbox'] = [xmin,ymin,w,h]
bbox_dict['segmentation'] = [[xmin,ymin,xmax,ymin,xmax,ymax,xmin,ymax]]
write_json_context['annotations'].append(bbox_dict)
num_bboxes+=1
i+=1
# 终于搞定了,保存！
with open(coco_format_save_path,'w') as fw:
json.dump(write_json_context,fw)
