modularization/src/detection/drivable_segmentation-master/generate_dataset.py

import os.path as ops
import json
import cv2
import numpy as np
import os
import glob
import shutil
def process_json_file(json_root,image_root,des_root):
    color_map = {"drivable": 1, "alternative": 2}
    if not os.path.exists(des_root):
        os.makedirs(des_root)
    images = des_root + '/images'
    labels = des_root + '/labels'
    if not os.path.exists(images):
        os.makedirs(images)
    if not os.path.exists(labels):
        os.makedirs(labels)
    i =0
    for js in glob.glob(json_root+'/**/*.json', recursive=True):
        path = js.split(json_root)[1][1:-5] # 4-5000\7\5183.png
        img_path = image_root+'/'+path
        rename = (path.replace('\\','-')).replace('/', '-') # 4-5000-7-5183.png
        with open(js, 'r') as f:
            dict = json.load(f)
            area_list = dict["annotations"]["area"]
        map = np.zeros([1080, 1920], dtype=np.uint8)
        for area in area_list:
            if area["points"]:
                a = [area["points"]]
            else:
                continue
            cv2.fillPoly(map, np.array(a).astype(int), color_map[area["category"]])
        label_name = rename.replace('.png', '_drivable.png')  # 4-5000-7-5183_drivable.png
        label_dir = des_root + '/' + 'labels'
        if not os.path.exists(label_dir):
            os.makedirs(label_dir)
        image_dir = des_root + '/' + 'images'
        if not os.path.exists(image_dir):
            os.makedirs(image_dir)
        label_path = label_dir + '/' + label_name
        image_path = image_dir + '/' + rename
        #resize_map = cv2.resize(map, (1280, 720), interpolation=cv2.INTER_AREA)
        cv2.imwrite(label_path, map)
        img = cv2.imread(img_path)
        #resize_img = cv2.resize(img, (1280,720), interpolation=cv2.INTER_AREA)
        cv2.imwrite(image_path,img)
        i+=1
        if i%500==0:
            print('完成'+str(i)+'张')
import random
def generate_train_val_test(des_root):
    images_dir=des_root+'/images'
    labels_dir = des_root + '/labels'
    images_list = os.listdir(images_dir)
    #split source file into 0.9, 0.05 , 0.05 for training, validation and testing
    random.shuffle(images_list)
    size = len(images_list)
    train = images_list[:int(size*0.9)]
    val = images_list[int(size * 0.9):int(size * 0.95)]
    test = images_list[int(size * 0.95):]
    random.shuffle(train)
    random.shuffle(val)
    random.shuffle(test)
    train_img_dir = des_root + '/dataset/images/train'
    val_img_dir = des_root + '/dataset/images/val'
    test_img_dir = des_root + '/dataset/images/test'
    train_lbl_dir = des_root + '/dataset/labels/train'
    val_lbl_dir = des_root + '/dataset/labels/val'
    test_lbl_dir = des_root + '/dataset/labels/test'
    if not os.path.exists(train_img_dir):
        os.makedirs(train_img_dir)
    if not os.path.exists(val_img_dir):
        os.makedirs(val_img_dir)
    if not os.path.exists(test_img_dir):
        os.makedirs(test_img_dir)
    if not os.path.exists(train_lbl_dir):
        os.makedirs(train_lbl_dir)
    if not os.path.exists(val_lbl_dir):
        os.makedirs(val_lbl_dir)
    if not os.path.exists(test_lbl_dir):
        os.makedirs(test_lbl_dir)
    for img in train:
        src_path = images_dir+'/'+img
        des_path = train_img_dir+'/'+img
        shutil.copy(src_path, des_path)
        lbl = img.replace('.png','_drivable.png')
        src_path = labels_dir+'/'+lbl
        des_path = train_lbl_dir+'/'+lbl
        shutil.copy(src_path, des_path)
    print('train done')
    for img in val:
        src_path = images_dir+'/'+img
        des_path = val_img_dir+'/'+img
        shutil.copy(src_path, des_path)
        lbl = img.replace('.png', '_drivable.png')
        src_path = labels_dir+'/'+lbl
        des_path = val_lbl_dir+'/'+lbl
        shutil.copy(src_path, des_path)
    print('val done')
    for img in test:
        src_path = images_dir+'/'+img
        des_path = test_img_dir+'/'+img
        shutil.copy(src_path, des_path)
        lbl = img.replace('.png', '_drivable.png')
        src_path = labels_dir+'/'+lbl
        des_path = test_lbl_dir+'/'+lbl
        shutil.copy(src_path, des_path)
    print('test done')

if __name__ == '__main__':
    json_root = '/media/adc/Elements/drivable/drivable_json'
    des_root = '/media/adc/Elements/drivable_dataset'
    image_root = '/media/adc/Elements/drivable/images'
    #合并并生成数据集图片
    process_json_file(json_root,image_root,des_root)
    #随机分配train\val\test
    generate_train_val_test(des_root)