目标检测 1:制作数据集

转自: 目标检测系列一：如何制作数据集？

txt to coco

# -*- coding: UTF-8 -*-
import cv2
import json
import sys


# process bar
def process_bar(count, total, status=''):
    bar_len = 60
    filled_len = int(round(bar_len * count / float(total)))

    percents = round(100.0 * count / float(total), 1)
    bar = '=' * filled_len + '-' * (bar_len - filled_len)

    sys.stdout.write('[%s] %s%s ...%s\r' % (bar, percents, '%', status))
    sys.stdout.flush()


root_path = "/Users/zcunyi/zcunyi/mot/val/"
images, categories, annotations = [], [], []

category_dict = {"people": 1}

for cat_n in category_dict:
    categories.append({"supercategory": "", "id": category_dict[cat_n], "name": cat_n})

# with open("train_clean.txt", "r") as f:
with open('val_clean.txt', 'r') as f:
    img_id = 0
    anno_id_count = 0
    count = 1
    # total = len(f.readlines())
    # print(111,total)
    lines = f.readlines()
    total = len(lines)
    # print(333,lines)
    for line in lines:
        # print(2222,line)
        # process_bar(count, total)

        print('{}/{}'.format(count,total))
        count += 1
        line = line.split(' ')
        img_name = line[0]
        bbox_num = int(line[1])
        img_cv2 = cv2.imread(root_path + img_name)
        [height, width, _] = img_cv2.shape
        bbox = line[1:]
        x = bbox[1::5]
        y = bbox[2::5]
        w = bbox[3::5]
        h = bbox[4::5]
        # img_name+=' '
        # fp.write(img_name)
        # for i in range(len(x)):

        # images info
        images.append({"file_name": img_name, "height": height, "width": width, "id": img_id})

        """
        annotation info:
        id : anno_id_count
        category_id : category_id
        bbox : bbox
        segmentation : [segment]
        area : area
        iscrowd : 0
        image_id : image_id
        """
        category_id = category_dict["people"]
        for i in range(0, len(x)):
            x1, y1, weight, height = int(x[i]), int(y[i]), int(w[i]), int(h[i])
            x2,y2 = x1+weight,y1+height
            # box = '{} {},{},{},{}\n'.format(img_name, x1, y1, x2, y2)
            # x1 = float(line[i * 5 + 3])
            # y1 = float(line[i * 5 + 4])
            # x2 = float(line[i * 5 + 3]) + float(line[i * 5 + 5])
            # y2 = float(line[i * 5 + 4]) + float(line[i * 5 + 6])
            # width = float(line[i * 5 + 5])
            # height = float(line[i * 5 + 6])

            bbox = [x1, y1, width, height]
            segment = [x1, y1, x2, y1, x2, y2, x1, y2]
            area = width * height

            anno_info = {'id': anno_id_count, 'category_id': category_id, 'bbox': bbox, 'segmentation': [segment],
                         'area': area, 'iscrowd': 0, 'image_id': img_id}
            annotations.append(anno_info)
            anno_id_count += 1

        img_id = img_id + 1

all_json = {"images": images, "annotations": annotations, "categories": categories}

with open("val.json", "w") as outfile:
    json.dump(all_json, outfile)

 

csv to coco

import os
import json
import numpy as np
import pandas as pd
import glob
import cv2
import os
import shutil
from IPython import embed
from sklearn.model_selection import train_test_split
np.random.seed(41)

#0为背景
# classname_to_id = {"person": 1}
classname_to_id = {"a":1, "b":2, "c":3, "d":4, "e":5, "f":6, "g":7, "h":8, "i":9, "j":10, "k":11, "l":12, "m":13, "n":14, "o":15, "p":16, "q":17, "r":18, "s":19, "t":20, "u":21}

class Csv2CoCo:

    def __init__(self,image_dir,total_annos):
        self.images = []
        self.annotations = []
        self.categories = []
        self.img_id = 0
        self.ann_id = 0
        self.image_dir = image_dir
        self.total_annos = total_annos

    def save_coco_json(self, instance, save_path):
        json.dump(instance, open(save_path, 'w'), ensure_ascii=False, indent=2)  # indent=2 更加美观显示
        # ensure_ascii=False 遇到中文会输出
        # dumps是将dict转化成str格式，loads是将str转化成dict格式;json_str = json.dumps(dataset,indent=2),f.write(json_str);
        # dump需要指定保存的文件对象

    # 由txt文件构建COCO
    def to_coco(self, keys):
        self._init_categories()
        number = 1
        for key in keys:
            print("~~~~~~~~~~{}/{}~~~~~~~~~~".format((number),len(keys)))
            number+=1
            self.images.append(self._image(key))
            shapes = self.total_annos[key]
            for shape in shapes:
                bboxi = []
                for cor in shape[:-1]:
                    bboxi.append(int(cor))
                label = shape[-1]
                annotation = self._annotation(bboxi,label)
                self.annotations.append(annotation)
                self.ann_id += 1
            self.img_id += 1
        instance = {}
        instance['info'] = 'spytensor created'
        instance['license'] = ['license']
        instance['images'] = self.images
        instance['annotations'] = self.annotations
        instance['categories'] = self.categories
        return instance

    # 构建类别
    def _init_categories(self):
        for k, v in classname_to_id.items():
            category = {}
            category['id'] = v
            category['name'] = k
            self.categories.append(category)

    # 构建COCO的image字段
    def _image(self, path):
        image = {}
        # print(path)
        img = cv2.imread(self.image_dir + path)
        image['height'] = img.shape[0]
        image['width'] = img.shape[1]
        image['id'] = self.img_id
        image['file_name'] = path
        return image

    # 构建COCO的annotation字段
    def _annotation(self, shape,label):
        # label = shape[-1]
        points = shape[:4] #看具体的csv文件,有的给出的是4个坐标,有的是8个坐标;[:4]针对4个坐标
        annotation = {}
        annotation['id'] = self.ann_id
        annotation['image_id'] = self.img_id
        # print(classname_to_id[label])
        # annotation['category_id'] = int(classname_to_id[label])
        annotation['category_id'] = int(label)
        annotation['segmentation'] = self._get_seg(points)
        annotation['bbox'] = self._get_box(points)
        annotation['iscrowd'] = 0
        annotation['area'] = 1.0
        return annotation

    # COCO的格式： [x1,y1,w,h] 对应COCO的bbox格式
    def _get_box(self, points):
        min_x = points[0]
        min_y = points[1]
        max_x = points[2]
        max_y = points[3]
        return [min_x, min_y, max_x - min_x, max_y - min_y]
    # segmentation
    def _get_seg(self, points):
        min_x = points[0]
        min_y = points[1]
        max_x = points[2]
        max_y = points[3]
        h = max_y - min_y
        w = max_x - min_x
        a = []
        a.append([min_x,min_y, min_x,min_y+0.5*h, min_x,max_y, min_x+0.5*w,max_y, max_x,max_y, max_x,max_y-0.5*h, max_x,min_y, max_x-0.5*w,min_y])
        return a
   

if __name__ == '__main__':
    csv_file = "train_label.csv"
    image_dir = "images/"
    saved_coco_path = "./"
    # 整合csv格式标注文件
    total_csv_annotations = {}
    annotations = pd.read_csv(csv_file,header=0).values #header=None,会自动加上列索引[0,1,2,3,,,],如果csv文件本身有列索引的话,设置header=0,以文件第一行为列索引;values取得是数组值
    for annotation in annotations:
        key = annotation[0].split(os.sep)[-1] #annotation[0] 一般是图片路径,data/path/to/img,用'os.sep'进行切分,[-1]取最后一个就是文件名:'img.jpg'
        value = np.array([annotation[1:]])
        if key in total_csv_annotations.keys(): #一张图中有多个object目标的情况
            total_csv_annotations[key] = np.concatenate((total_csv_annotations[key],value),axis=0)
        else:
            total_csv_annotations[key] = value
    # 按照键值划分数据
    total_keys = list(total_csv_annotations.keys())
    train_keys, val_keys = train_test_split(total_keys, test_size=0.2)
    print("train_n:", len(train_keys), 'val_n:', len(val_keys))
    # 创建必须的文件夹
    if not os.path.exists('%scoco/annotations/'%saved_coco_path):
        os.makedirs('%scoco/annotations/'%saved_coco_path)
    if not os.path.exists('%scoco/images/train2017/'%saved_coco_path):
        os.makedirs('%scoco/images/train2017/'%saved_coco_path)
    if not os.path.exists('%scoco/images/val2017/'%saved_coco_path):
        os.makedirs('%scoco/images/val2017/'%saved_coco_path)
    # 把训练集转化为COCO的json格式
    l2c_train = Csv2CoCo(image_dir=image_dir,total_annos=total_csv_annotations)
    train_instance = l2c_train.to_coco(train_keys)
    l2c_train.save_coco_json(train_instance, '%scoco/annotations/instances_train2017.json'%saved_coco_path)
    for file in train_keys:
        shutil.copy(image_dir+file,"%scoco/images/train2017/"%saved_coco_path) #shutil.copy(源文件,目的文件夹)
    for file in val_keys:
        shutil.copy(image_dir+file,"%scoco/images/val2017/"%saved_coco_path)
    # 把验证集转化为COCO的json格式
    l2c_val = Csv2CoCo(image_dir=image_dir,total_annos=total_csv_annotations)
    val_instance = l2c_val.to_coco(val_keys)
    l2c_val.save_coco_json(val_instance, '%scoco/annotations/instances_val2017.json'%saved_coco_path)

csv to voc

import os
import numpy as np
import codecs
import pandas as pd
import json
from glob import glob
import cv2
import shutil
from sklearn.model_selection import train_test_split
from IPython import embed
#1.标签路径
# csv_file = "../csv/train_labels.csv"
csv_file = "train_label.csv"
saved_path = "./VOCdevkit/VOC2007/"                #保存路径
image_save_path = "./JPEGImages/"
image_raw_parh = "images/"
#2.创建要求文件夹
if not os.path.exists(saved_path + "Annotations"):
    os.makedirs(saved_path + "Annotations")
if not os.path.exists(saved_path + "JPEGImages/"):
    os.makedirs(saved_path + "JPEGImages/")
if not os.path.exists(saved_path + "ImageSets/Main/"):
    os.makedirs(saved_path + "ImageSets/Main/")
    
#3.获取待处理文件
total_csv_annotations = {}
annotations = pd.read_csv(csv_file,header=0).values
for annotation in annotations:
    key = annotation[0].split(os.sep)[-1]
    value = np.array([annotation[1:]])
    if key in total_csv_annotations.keys():
        total_csv_annotations[key] = np.concatenate((total_csv_annotations[key],value),axis=0)
    else:
        total_csv_annotations[key] = value
# print("1111",total_csv_annotations)
#4.读取标注信息并写入 xml
for filename,label in total_csv_annotations.items():
    #embed()
    print("222",label)
    height, width, channels = cv2.imread(image_raw_parh + filename).shape #HWC
    #embed()
    with codecs.open(saved_path + "Annotations/"+filename.replace(".jpg",".xml"),"w","utf-8") as xml:
        xml.write('<annotation>\n')
        xml.write('\t<folder>' + 'UAV_data' + '</folder>\n')
        xml.write('\t<filename>' + filename + '</filename>\n')
        xml.write('\t<source>\n')
        xml.write('\t\t<database>The UAV autolanding</database>\n')
        xml.write('\t\t<annotation>UAV AutoLanding</annotation>\n')
        xml.write('\t\t<image>flickr</image>\n')
        xml.write('\t\t<flickrid>NULL</flickrid>\n')
        xml.write('\t</source>\n')
        xml.write('\t<owner>\n')
        xml.write('\t\t<flickrid>NULL</flickrid>\n')
        xml.write('\t\t<name>ChaojieZhu</name>\n')
        xml.write('\t</owner>\n')
        xml.write('\t<size>\n')
        xml.write('\t\t<width>'+ str(width) + '</width>\n')
        xml.write('\t\t<height>'+ str(height) + '</height>\n')
        xml.write('\t\t<depth>' + str(channels) + '</depth>\n')
        xml.write('\t</size>\n')
        xml.write('\t\t<segmented>0</segmented>\n')
        if isinstance(label,float):
            ## 空白
            xml.write('</annotation>')
            continue
        for label_detail in label:
            labels = label_detail
            #看具体的csv文件,有的给出的是4个坐标,有的是8个坐标;
            xmin = int(labels[0])
            ymin = int(labels[1])
            xmax = int(labels[4])
            ymax = int(labels[5])
            label_ = labels[-1] #有的会是label名,str类型;有的为int类型;int类型,需要在下面xml.write的时候加上'str'
            if xmax <= xmin:
                pass
            elif ymax <= ymin:
                pass
            else:
                xml.write('\t<object>\n')
                xml.write('\t\t<name>'+str(label_)+'</name>\n')
                xml.write('\t\t<pose>Unspecified</pose>\n')
                xml.write('\t\t<truncated>1</truncated>\n')
                xml.write('\t\t<difficult>0</difficult>\n')
                xml.write('\t\t<bndbox>\n')
                xml.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')
                xml.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')
                xml.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')
                xml.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')
                xml.write('\t\t</bndbox>\n')
                xml.write('\t</object>\n')
                print(filename,xmin,ymin,xmax,ymax,labels)
        xml.write('</annotation>')
        

#6.split files for txt
txtsavepath = saved_path + "ImageSets/Main/"
ftrainval = open(txtsavepath+'/trainval.txt', 'w')
ftest = open(txtsavepath+'/test.txt', 'w')
ftrain = open(txtsavepath+'/train.txt', 'w')
fval = open(txtsavepath+'/val.txt', 'w')
total_files = glob(saved_path+"./Annotations/*.xml")
total_files = [i.split("/")[-1].split(".xml")[0] for i in total_files]
#test_filepath = ""
for file in total_files:
    ftrainval.write(file + "\n")

# move images to voc JPEGImages folder
for image in glob(image_raw_parh+"/*.jpg"):
    shutil.copy(image,saved_path+image_save_path)

train_files,val_files = train_test_split(total_files,test_size=0.15,random_state=42)

for file in train_files:
    ftrain.write(file + "\n")
#val
for file in val_files:
    fval.write(file + "\n")

ftrainval.close()
ftrain.close()
fval.close()
#ftest.close()

csv to labelme

import os
import cv2
import json
import pandas as pd
import numpy as np
from glob import glob 
from tqdm import tqdm
from IPython import embed
import base64
# from labelme import utils # 激活labelme虚拟环境
image_path = "./images/"
csv_file = "./train_label.csv"
annotations = pd.read_csv(csv_file,header=0).values
total_csv_annotations = {}
for annotation in annotations:
    key = annotation[0].split(os.sep)[-1]
    value = np.array([annotation[1:]])
    if key in total_csv_annotations.keys():
        total_csv_annotations[key] = np.concatenate((total_csv_annotations[key],value),axis=0)
    else:
        total_csv_annotations[key] = value
for key,value in total_csv_annotations.items():
    height,width = cv2.imread(image_path+key).shape[:2]
    labelme_format = {
    "version":"3.6.16",
    "flags":{},
    "lineColor":[0,255,0,128],
    "fillColor":[255,0,0,128],
    "imagePath":key,
    "imageHeight":height,
    "imageWidth":width
    }
    with open(image_path+key,"rb") as f:
        imageData = f.read()
        imageData = base64.b64encode(imageData).decode('utf-8')
    #img = utils.img_b64_to_arr(imageData)
    labelme_format["imageData"] = imageData
    shapes = []
    for shape in value:
        label = shape[-1]
        s = {"label":label,"line_color":None,"fill_color":None,"shape_type":"rectangle"}
        points = [
            [shape[0],shape[1]],
            [shape[2],shape[3]]
            ]
        s["points"] = points
        shapes.append(s)
    labelme_format["shapes"] = shapes
    json.dump(labelme_format,open("%s%s"%(image_path,key.replace(".jpg",".json")),"w"),ensure_ascii=False, indent=2)

labelme to coco

import os
import json
import numpy as np
import glob
import shutil
from sklearn.model_selection import train_test_split
np.random.seed(41)

#0为背景
classname_to_id = {"person": 1}

class Lableme2CoCo:

    def __init__(self):
        self.images = []
        self.annotations = []
        self.categories = []
        self.img_id = 0
        self.ann_id = 0

    def save_coco_json(self, instance, save_path):
        json.dump(instance, open(save_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)  # indent=2 更加美观显示

    # 由json文件构建COCO
    def to_coco(self, json_path_list):
        self._init_categories()
        for json_path in json_path_list:
            obj = self.read_jsonfile(json_path)
            self.images.append(self._image(obj, json_path))
            shapes = obj['shapes']
            for shape in shapes:
                annotation = self._annotation(shape)
                self.annotations.append(annotation)
                self.ann_id += 1
            self.img_id += 1
        instance = {}
        instance['info'] = 'spytensor created'
        instance['license'] = ['license']
        instance['images'] = self.images
        instance['annotations'] = self.annotations
        instance['categories'] = self.categories
        return instance

    # 构建类别
    def _init_categories(self):
        for k, v in classname_to_id.items():
            category = {}
            category['id'] = v
            category['name'] = k
            self.categories.append(category)

    # 构建COCO的image字段
    def _image(self, obj, path):
        image = {}
        from labelme import utils
        img_x = utils.img_b64_to_arr(obj['imageData'])
        h, w = img_x.shape[:-1]
        image['height'] = h
        image['width'] = w
        image['id'] = self.img_id
        image['file_name'] = os.path.basename(path).replace(".json", ".jpg")
        return image

    # 构建COCO的annotation字段
    def _annotation(self, shape):
        label = shape['label']
        points = shape['points']
        annotation = {}
        annotation['id'] = self.ann_id
        annotation['image_id'] = self.img_id
        annotation['category_id'] = int(classname_to_id[label])
        annotation['segmentation'] = [np.asarray(points).flatten().tolist()]
        annotation['bbox'] = self._get_box(points)
        annotation['iscrowd'] = 0
        annotation['area'] = 1.0
        return annotation

    # 读取json文件，返回一个json对象
    def read_jsonfile(self, path):
        with open(path, "r", encoding='utf-8') as f:
            return json.load(f)

    # COCO的格式： [x1,y1,w,h] 对应COCO的bbox格式
    def _get_box(self, points):
        min_x = min_y = np.inf
        max_x = max_y = 0
        for x, y in points:
            min_x = min(min_x, x)
            min_y = min(min_y, y)
            max_x = max(max_x, x)
            max_y = max(max_y, y)
        return [min_x, min_y, max_x - min_x, max_y - min_y]


if __name__ == '__main__':
    labelme_path = "labelme/"
    saved_coco_path = "./"
    # 创建文件
    if not os.path.exists("%scoco/annotations/"%saved_coco_path):
        os.makedirs("%scoco/annotations/"%saved_coco_path)
    if not os.path.exists("%scoco/images/train2017/"%saved_coco_path):
        os.makedirs("%scoco/images/train2017"%saved_coco_path)
    if not os.path.exists("%scoco/images/val2017/"%saved_coco_path):
        os.makedirs("%scoco/images/val2017"%saved_coco_path)
    # 获取images目录下所有的joson文件列表
    json_list_path = glob.glob(labelme_path + "/*.json")
    # 数据划分,这里没有区分val2017和tran2017目录，所有图片都放在images目录下
    train_path, val_path = train_test_split(json_list_path, test_size=0.12)
    print("train_n:", len(train_path), 'val_n:', len(val_path))

    # 把训练集转化为COCO的json格式
    l2c_train = Lableme2CoCo()
    train_instance = l2c_train.to_coco(train_path)
    l2c_train.save_coco_json(train_instance, '%scoco/annotations/instances_train2017.json'%saved_coco_path)
    for file in train_path:
        shutil.copy(file.replace("json","jpg"),"%scoco/images/train2017/"%saved_coco_path)
    for file in val_path:
        shutil.copy(file.replace("json","jpg"),"%scoco/images/val2017/"%saved_coco_path)

    # 把验证集转化为COCO的json格式
    l2c_val = Lableme2CoCo()
    val_instance = l2c_val.to_coco(val_path)
    l2c_val.save_coco_json(val_instance, '%scoco/annotations/instances_val2017.json'%saved_coco_path)

 

labelme to voc

import os
import numpy as np
import codecs
import json
from glob import glob
import cv2
import shutil
from sklearn.model_selection import train_test_split
#1.标签路径
labelme_path = "./labelme/"              #原始labelme标注数据路径
saved_path = "./VOCdevkit/VOC2019/"                #保存路径

#2.创建要求文件夹
if not os.path.exists(saved_path + "Annotations"):
    os.makedirs(saved_path + "Annotations")
if not os.path.exists(saved_path + "JPEGImages/"):
    os.makedirs(saved_path + "JPEGImages/")
if not os.path.exists(saved_path + "ImageSets/Main/"):
    os.makedirs(saved_path + "ImageSets/Main/")
    
#3.获取待处理文件
files = glob(labelme_path + "*.json")
files = [i.split("/")[-1].split(".json")[0] for i in files]

#4.读取标注信息并写入 xml
for json_file_ in files:
    json_filename = labelme_path + json_file_ + ".json"
    json_file = json.load(open(json_filename,"r",encoding="utf-8"))
    height, width, channels = cv2.imread(labelme_path + json_file_ +".jpg").shape
    with codecs.open(saved_path + "Annotations/"+json_file_ + ".xml","w","utf-8") as xml:
        xml.write('<annotation>\n')
        xml.write('\t<folder>' + 'UAV_data' + '</folder>\n')
        xml.write('\t<filename>' + json_file_ + ".jpg" + '</filename>\n')
        xml.write('\t<source>\n')
        xml.write('\t\t<database>The UAV autolanding</database>\n')
        xml.write('\t\t<annotation>UAV AutoLanding</annotation>\n')
        xml.write('\t\t<image>flickr</image>\n')
        xml.write('\t\t<flickrid>NULL</flickrid>\n')
        xml.write('\t</source>\n')
        xml.write('\t<owner>\n')
        xml.write('\t\t<flickrid>NULL</flickrid>\n')
        xml.write('\t\t<name>ChaojieZhu</name>\n')
        xml.write('\t</owner>\n')
        xml.write('\t<size>\n')
        xml.write('\t\t<width>'+ str(width) + '</width>\n')
        xml.write('\t\t<height>'+ str(height) + '</height>\n')
        xml.write('\t\t<depth>' + str(channels) + '</depth>\n')
        xml.write('\t</size>\n')
        xml.write('\t\t<segmented>0</segmented>\n')
        for multi in json_file["shapes"]:
            points = np.array(multi["points"])
            # labelme标注的可能是多边形,记录的是每个点的x,y坐标,对所有x,y坐标求取min和max,即可得出目标的bbox.
            xmin = min(points[:,0])
            xmax = max(points[:,0])
            ymin = min(points[:,1])
            ymax = max(points[:,1])
            label = multi["label"]
            if xmax <= xmin:
                pass
            elif ymax <= ymin:
                pass
            else:
                xml.write('\t<object>\n')
                xml.write('\t\t<name>'+label+'</name>\n')
                xml.write('\t\t<pose>Unspecified</pose>\n')
                xml.write('\t\t<truncated>1</truncated>\n')
                xml.write('\t\t<difficult>0</difficult>\n')
                xml.write('\t\t<bndbox>\n')
                xml.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')
                xml.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')
                xml.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')
                xml.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')
                xml.write('\t\t</bndbox>\n')
                xml.write('\t</object>\n')
                print(json_filename,xmin,ymin,xmax,ymax,label)
        xml.write('</annotation>')
        
#5.复制图片到 VOC2007/JPEGImages/下
image_files = glob(labelme_path + "*.jpg")
print("copy image files to VOC019/JPEGImages/")
for image in image_files:
    shutil.copy(image,saved_path +"JPEGImages/")
    
#6.split files for txt
txtsavepath = saved_path + "ImageSets/Main/"
ftrainval = open(txtsavepath+'/trainval.txt', 'w')
ftest = open(txtsavepath+'/test.txt', 'w')
ftrain = open(txtsavepath+'/train.txt', 'w')
fval = open(txtsavepath+'/val.txt', 'w')
total_files = glob("./VOCdevkit/VOC2019/Annotations/*.xml")
total_files = [i.split("/")[-1].split(".xml")[0] for i in total_files]
#test_filepath = ""
for file in total_files:
    ftrainval.write(file + "\n")
#test
#for file in os.listdir(test_filepath):
#    ftest.write(file.split(".jpg")[0] + "\n")
#split
train_files,val_files = train_test_split(total_files,test_size=0.15,random_state=42)
#train
for file in train_files:
    ftrain.write(file + "\n")
#val
for file in val_files:
    fval.write(file + "\n")

ftrainval.close()
ftrain.close()
fval.close()
#ftest.close()

csv to coco_train

# -*- coding: utf-8 -*-
import os,shutil
import json
import cv2
import platform
# class_all=['0','1','2','3','4','5','6','7','8','9','10','11','12','13','14','15','16','17','18','19','20']
class_all = ['Z','A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T']
if platform.platform()[:5]=='Linux':
    # train_dir='./coco_traffic_sign_recognition_train_dataset'
    # train_dir='/home/zcy/anaconda3/envs/detectron/detectron/detectron/datasets/data/coco_traffic_sign_recognition/coco_traffic_sign_recognition_train_dataset/'
    train_dir = '/home/zcy/0_projects/prepare_detection_dataset/images/'
    train_csv_file='/home/zcy/0_projects/prepare_detection_dataset/train_label.csv'
    json_file = '001.json'
    # json_file = 'coco_traffic_sign_recognition_train_new_cut.json'
    # json_file = 'aa.json'
# if not os.path.exists('data/annotations'):
#     os.mkdir('data/annotations')
else:
    name="data_augmentation"
    train_dir='G:\\MyFiles\Lab542\\rebar\\detectron\\'+name+'\\img\\'
    train_csv_file='G:\\MyFiles\Lab542\\rebar\\detectron\\'+name+'\\gt.csv'
    json_file = 'G:\\MyFiles\Lab542\\rebar\\detectron\\'+name+'\\df_train_da.json'

# if not os.path.exists('data/annotations'):
#     os.mkdir('data/annotations')
load_file_list=False
def safe2coco():
    if load_file_list:
        with open(json_file) as f:
            dataset=json.load(f)
        dataset['annotations']=[]
    else:
        dataset = {
        'licenses': [],
        'info': {},
        'categories': [],
        'images': [],
        'annotations': []
        }
        for i, cls in enumerate(class_all,1):
            dataset['categories'].append({
              'id': i,
              'name': cls,
              'supercategory': 'traffic_sign'
          })
        file_id=0
        file_list=os.listdir(train_dir)
        # print(train_dir)
        count = len([name for name in os.listdir(train_dir) if os.path.isfile(os.path.join(train_dir, name))])
        # print("~~~count~~~: ", count)
        for file in file_list:
            # print(train_dir + file)
            img = cv2.imread(train_dir+file)
            width=img.shape[1]
            height=img.shape[0]
            dataset['images'].append({
                  'coco_url': '',
                  'date_captured': '',
                  'file_name': file,
                  'flickr_url': '',
                  'id': file_id,
                  'license': 0,
                  'width': width,
                  'height': height
                  })
            # print(file,width,height)
            file_id+=1
            print("~~~~~~~~~~{}/{}~~~~~~~~~~".format((file_id),count))

    file_dict={}
    for file in dataset['images']:
        file_dict[file['file_name']]=file['id']

    box_id=0
    with open(train_csv_file) as f:
        lines=f.readlines()
    for line in lines[1:]:
        line=line.replace('\n','')
        split=line.split(',')
        # split=line.split(' ')
        filename=split[0]
        xmin = int(split[1])
        ymin = int(split[2])
        xmax = int(split[5])
        ymax = int(split[6])
        w = xmax-xmin
        h = ymax-ymin
        file_id = file_dict[filename]
        if w * h > 1*1:
            dataset['annotations'].append({
              'area': w * h,
              'bbox': [xmin, ymin, w, h],
              # 'category_id': 1,
              'category_id': int(split[-1]),
              'id': box_id,
              'image_id': file_id,
              'iscrowd': 0,
              'segmentation': []
            })
            box_id+=1
        else:
            print('skip',filename)
        # box=split[1]
        # if box!='':
        #     box=box.split(' ')
        #     xmin=int(box[0])
        #     ymin=int(box[1])
        #     xmax=int(box[2])
        #     ymax=int(box[3])
        #     w=xmax-xmin+1
        #     h=ymax-ymin+1
        #     file_id=file_dict[filename]
        #     if w*h>16*16:
        #         dataset['annotations'].append({
        #           'area': w * h,
        #           'bbox': [xmin, ymin, w, h],
        #           'category_id': 1,
        #           'id': box_id,
        #           'image_id': file_id,
        #           'iscrowd': 0,
        #           'segmentation': []
        #         })
        #         box_id+=1
        #     else:
        #         print('skip',box,filename)

            #print(filename,file_dict[filename],box)
    with open(json_file, 'w') as f:
        #json.dump(dataset, f)
        json_str = json.dumps(dataset,indent=2)
        f.write(json_str)
    print("------------------json write end   111-------------------")
def json_pop_image_without_box():
    with open(json_file) as f:
        dataset=json.load(f)
    print('total images:',len(dataset['images']))
    print('total boxes:',len(dataset['annotations']))
    file_id_list=[]
    for box in dataset['annotations']:
        if box['image_id'] not in file_id_list:
            file_id_list.append(box['image_id'])
    print('file with box:',len(file_id_list))
    dataset['images']=[dataset['images'][x] for x in file_id_list]
    print('after filter,image count:',len(dataset['images']))
    with open(json_file.replace('.json','_filter.json'), 'w') as f:
        #json.dump(dataset, f)
        json_str = json.dumps(dataset)
        f.write(json_str)
    print("------------------json write end   222-------------------")
def print_json_info(json_file):
    # area_thresh=[0,8,16,32,64,128,256,512,1024,2048]
    # area_count=[0,0,0,0,0,0,0,0,0,0]
    area_thresh = [0,10,20,30,40,50,60,70,80,90,100,120,128,180,500]
    area_count = [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
    h_w = [0.8,0.85,0.88,0.9,0.95,0.98,1.0,1.05,1.1,1.15,1.2,1.25,1.3,2.0]
    h_w_count = [0,0,0,0,0,0,0,0,0,0,0,0,0,0]
    with open(json_file) as f:
        dataset=json.load(f)
    print('total images:',len(dataset['images']))
    print('total boxes:',len(dataset['annotations']))
    file_id_list=[]
    for box in dataset['annotations']:
        # print(box['bbox'])
        if box['image_id'] not in file_id_list:
            file_id_list.append(box['image_id'])
        for index,area in enumerate(area_thresh):
            if box['area']<=(area*area):
                area_count[index]+=1
                break
    print('file with box:',len(file_id_list))
    print('area thresh    :',area_thresh)
    print('area statistics:',area_count)
    for box in dataset['annotations']:
        for index, hw in enumerate(h_w):
            if box['bbox'][3]/box['bbox'][2] <= hw*hw:
                h_w_count[index] += 1
                break
    print("h/w score :", h_w)
    print("h/w thresh :", h_w_count)
    for box in dataset['annotations']:
        if box['bbox'][3] / box['bbox'][2] >2.0 :
            print()
if __name__=='__main__':
    #safe2coco()
    safe2coco()
    #json_pop_image_without_box()
    # print_json_info(json_file)