SAR/crop.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 crop(src_path="d:/data/data_train_val_test",dst_path="d:/data/data_crop",win_width=640,win_height=640,step=320):
    if os.path.exists(dst_path):
        shutil.rmtree(dst_path)
    os.mkdir(dst_path)
    for phase in ["train","val"]:
        img_num=0
        if os.path.exists(dst_path+"/"+phase):
            shutil.rmtree(dst_path+"/"+phase)
        os.mkdir(dst_path+"/"+phase)
        f_phase=open(src_path+"/"+phase+".txt","r",encoding="utf-8")
        for line in f_phase.readlines():
            line=line.strip()
            ext = line.rsplit(".")[-1]
            len1 = len(line) - len(ext)
            json_f = line[:len1] + "json"  #json文件


            basename = os.path.basename(line)
            bsname = basename.rsplit(".", 1)[0]
            parent = os.path.dirname(line)

            if os.path.exists(json_f):
                f_json=open(json_f,"rb")
            else:
                json_f2 = os.path.join(parent, "cor_" + bsname + ".json")
                if os.path.exists(json_f2):
                    f_json=open(json_f2,"rb")
                    #print("crop json_f: ",json_f2)

            #f_json=open(json_f,"rb")
            json_dict = json.load(f_json)

            # if img_bsname in ["201406300029"]:
            #     continue
            img_ext=json_dict["imagePath"].rsplit(".",1)[-1]
            img=cv2.imread(line)
            big_h=json_dict["imageHeight"]
            big_w=json_dict["imageWidth"]



            for i in range(0,big_w,step):
                right_i = min(i + win_width, big_w)
                if right_i==i:
                    continue
                for j in range(0,big_h,step):
                    json_crop = {"version": "5.1.1", "flags": {}, "shapes": [], "imagePath": "", "imageHeight": 0,
                                 "imageWidth": 0,"imageData":""}
                    bottom_j=min(j+win_height,big_h)
                    if bottom_j == j:
                        continue
                    #crop_img=np.zeros([bottom_j-j,right_i-i,3])
                    #crop_img = img[j:bottom_j,i:right_i,:]
                    #按窗口大小裁剪小图，不够的地方用0填充。
                    crop_img = np.zeros([win_height, win_width, 3],img.dtype)
                    #print("j:bottom_j,i:right_i:", j,bottom_j,i,right_i,big_w,big_h)
                    crop_img[:bottom_j-j,:right_i-i,:] = img[j:bottom_j,i:right_i,:]
                    new_bsname=str(img_num)+"_"+str(i)+"_"+str(j)
                    json_crop["imagePath"]=new_bsname+"."+img_ext
                    json_crop["imageHeight"]=crop_img.shape[0]
                    json_crop["imageWidth"]=crop_img.shape[1]

                    dst_img_path=dst_path+"/"+phase+"/"+new_bsname+"."+img_ext
                    dst_json_path = dst_path+"/"+phase+"/" + new_bsname + ".json"


                    for shape in json_dict["shapes"]:
                        ps=[]
                        # for p in shape["points"]:
                        #     if p[0]>=i and p[0]<right_i and p[1]>=j and p[1]<bottom_j:
                        #         ps.append([p[0]-i,p[1]-j])

                        #
                        # if img_bsname in ["201406300015_0_0"]:
                        #     print("shape: ", shape)
                        #如果是圆，就将圆转化为多边形
                        if len(shape["points"])==2:
                            #print("circle,dst_img_path:",dst_img_path)
                            r=math.sqrt(math.pow(shape["points"][0][0]-shape["points"][1][0],2)+math.pow(shape["points"][0][1]-shape["points"][1][1],2))
                            circle = Point(shape["points"][0]).buffer(r)
                            try:
                                poly1=Polygon(circle)
                            except Exception as e:
                                print("poly1 circle error: ", e)
                        else:
                            try:
                                poly1=Polygon(shape["points"])
                            except Exception as e:
                                print("poly1 not circle error: ", e)
                        # except:
                        #     print("poly1=Polygon(shape[points]) error")
                        #     print("dst_img_path,shape[points]: ", dst_img_path, shape["points"])
                        #     continue
                        try:
                            poly2= box(i,j,right_i,bottom_j)
                        except Exception as e:
                            print("poly2 error: ", e)
                        try:
                            intersection = poly1.intersection(poly2)
                        except shapely.errors.GEOSException as e:
                            print("shapely.errors.GEOSException error:",e)
                            continue

                        if intersection.is_empty:
                            continue

                        #长形装备会被拆分为多个
                        if isinstance(intersection, MultiPolygon):
                            pss = []
                            #print("intersection type: ", intersection.type)
                            #print("intersection:", len(intersection.geoms))
                            for g in intersection.geoms:
                                if (g.area < 10):  #剔除面积太小的多边形
                                    continue
                                pss.append(list(g.exterior.coords))
                            for ps in pss:
                                ps_new = []
                                for p in ps:
                                    ps_new.append([p[0] - i, p[1] - j])
                                shape_dict = {"points": ps_new, "group_id": shape["group_id"],"shape_type": "polygon", "flags": {}}
                                if shape["label"] in ["white_crack","白色裂纹","白裂纹","白色裂缝","白裂缝","白色裂隙","白裂隙","balck_crack","black_crack","黑色裂纹","黑裂纹","黑色裂缝","黑裂缝","黑色裂隙","黑裂隙","crack","裂纹","裂缝","裂隙"]:
                                    label ="crack"
                                    shape_dict["label"]=label
                                elif shape["label"] in ["white_hole","白色孔洞","白孔洞","black_hole","黑色孔洞","黑孔洞","hole","孔洞"]:
                                    label ="hole"
                                    shape_dict["label"] = label
                                elif shape["label"] in ["white_debonding","白色脱粘","白脱粘","black_debonding","黑色脱粘","黑脱粘","debonding","脱粘"]:
                                    label ="debonding"
                                    shape_dict["label"] = label
                                elif shape["label"] in ["rarefaction","疏松"]:
                                    label ="rarefaction"
                                    shape_dict["label"] = label
                                else:
                                    pass
                                if "label" in list(shape_dict.keys()):
                                    json_crop["shapes"].append(shape_dict)

                        elif isinstance(intersection, Polygon):
                            if(intersection.area<10):
                                continue
                            ps=(list(intersection.exterior.coords))
                            ps_new = []
                            for p in ps:
                                ps_new.append([p[0] - i, p[1] - j])


                            shape_dict = {"points": ps_new, "group_id": shape["group_id"], "shape_type": "polygon",
                                          "flags": {}}
                            if shape["label"] in ["white_crack", "白色裂纹", "白裂纹", "白色裂缝", "白裂缝", "白色裂隙",
                                                  "白裂隙", "balck_crack", "black_crack", "黑色裂纹", "黑裂纹",
                                                  "黑色裂缝", "黑裂缝", "黑色裂隙", "黑裂隙", "crack", "裂纹", "裂缝",
                                                  "裂隙"]:
                                label = "crack"
                                shape_dict["label"] = label
                            elif shape["label"] in ["white_hole", "白色孔洞", "白孔洞", "black_hole", "黑色孔洞",
                                                    "黑孔洞", "hole", "孔洞"]:
                                label = "hole"
                                shape_dict["label"] = label
                            elif shape["label"] in ["white_debonding", "白色脱粘", "白脱粘", "black_debonding",
                                                    "黑色脱粘", "黑脱粘", "debonding", "脱粘"]:
                                label = "debonding"
                                shape_dict["label"] = label
                            elif shape["label"] in ["rarefaction", "疏松"]:
                                label = "rarefaction"
                                shape_dict["label"] = label
                            else:
                                pass
                            if "label" in list(shape_dict.keys()):
                                json_crop["shapes"].append(shape_dict)

                    if len(json_crop["shapes"])>0:   #只保存有目标的小图
                        cv2.imwrite(dst_img_path, crop_img)
                        def load_img(filename):
                            # pillow读取tif图片出错，所以只能用opencv读取，再转为pillow
                            # image_pil = PIL.Image.open(img_cv = cv2.imread(img_path)
                            img_cv = cv2.imread(filename)
                            image_pil = PIL.Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))
                            with io.BytesIO() as f:
                                ext = osp.splitext(filename)[1].lower()

                                if ext in [".jpg", ".jpeg"]:
                                    format = "JPEG"
                                else:
                                    format = "PNG"

                                image_pil.save(f, format=format)
                                f.seek(0)
                                return f.read()
                        imageData = base64.b64encode(load_img(dst_img_path)).decode("utf-8")
                        json_crop["imageData"] = imageData

                        jsonOrderedFile = json.dumps(json_crop, ensure_ascii=False, indent=2)
                        with open(dst_json_path, 'w') as jsonfile:
                            jsonfile.write(jsonOrderedFile)






            f_json.close()
            img_num += 1
        f_phase.close()