SAR/json_to_yolo.py

import glob
import json
import os
import shutil
import random
import cv2
import numpy as np
import base64
import math
import io
import PIL.Image
import os.path as osp

import shapely
from shapely.geometry import box, Polygon,MultiPolygon
from shapely.geometry import Point


def json_to_yolo(src_path="d:/data/data_crop",dst_path="d:/data/data_yolo"):
    if os.path.exists(dst_path):
        shutil.rmtree(dst_path)
    os.mkdir(dst_path)
    os.makedirs(os.path.join(dst_path,"images"),exist_ok=True)
    os.makedirs(os.path.join(dst_path, "labels"), exist_ok=True)
    #cls_names=["white_crack","white_hole","white_debonding","black_crack","black_hole","black_debonding","rarefaction","structure_circle","structure_smallrect","structure_longrect"]
    #cls_names=["white_crack","white_hole","white_debonding","black_crack","black_hole","black_debonding","rarefaction","circle_obj","rect_obj"]
    cls_names=["crack","hole","debonding","rarefaction"]
    for phase in ["train","val"]:
        if os.path.exists(dst_path+"/images/"+phase):
            shutil.rmtree(dst_path+"/images/"+phase)
        os.mkdir(dst_path+"/images/"+phase)

        if os.path.exists(dst_path+"/labels/"+phase):
            shutil.rmtree(dst_path+"/labels/"+phase)
        os.mkdir(dst_path+"/labels/"+phase)

        for file in glob.glob(src_path+"/"+phase+"/*"):

            bsname, ext = os.path.basename(file).rsplit(".", 1)
            if ext in ["tif","tiff", "png", "jpg", "jpeg","bmp"]:
                json_file = os.path.join(os.path.dirname(file),bsname+".json")
                #print("yolo json_file: ", json_file)
                if os.path.exists(json_file):
                    f_json=open(json_file,"rb")
                    json_dict = json.load(f_json)
                    #img_path = os.path.dirname(file)+"/"+json_dict["imagePath"]

                    #parent = os.path.basename(os.path.dirname(file))
                    shutil.copy(file,dst_path+"/images/"+phase+"/"+json_dict["imagePath"])

                    h=json_dict["imageHeight"]
                    w=json_dict["imageWidth"]
                    img_bsname = json_dict["imagePath"].rsplit(".", 1)[0]
                    #print(dst_path+"/labels/"+phase+"/"+img_bsname+".txt")
                    with open(dst_path+"/labels/"+phase+"/"+img_bsname+".txt","w") as f_yolo:
                        for shape in json_dict["shapes"]:
                            lines = ""
                            # if shape["label"] not in cls_names:
                            #     cls_names.append(shape["label"])
                            if shape["label"] not in cls_names:
                                continue
                            # if cls_names.index(shape["label"])>9:
                            #     print(shape['label'])
                            lines+=str(cls_names.index(shape["label"]))+" "
                            # if shape["label"]=="crack":
                            #     print(file)
                            for p in shape["points"]:
                                lines+=str(p[0]*1.0/w)+" "+str(p[1]*1.0/h)+" "

                            lines+="\n"
                            f_yolo.writelines(lines)
                elif phase in ["val"]:  #如果val中图片不存在对应的json文件，也加入到yolo的验证中
                    shutil.copy(file, dst_path + "/images/" + phase + "/" + bsname + "." + ext)
                    f_yolo = open(dst_path + "/labels/" + phase + "/" + bsname + ".txt", "w")
                    f_yolo.close()


        #print("cls_name: ",cls_names)