From fc660d61be7cd089c41d2882726b8fec487bcd31 Mon Sep 17 00:00:00 2001 From: Ali Samhat <113544685+AliSamhat07@users.noreply.github.com> Date: Sat, 16 Dec 2023 09:50:40 +0200 Subject: [PATCH] Add files via upload --- data_processing/generate_masks/Masker.py | 277 +++++++++++++++++++++++ 1 file changed, 277 insertions(+) create mode 100644 data_processing/generate_masks/Masker.py diff --git a/data_processing/generate_masks/Masker.py b/data_processing/generate_masks/Masker.py new file mode 100644 index 0000000..71717e8 --- /dev/null +++ b/data_processing/generate_masks/Masker.py @@ -0,0 +1,277 @@ +import numpy as np +import cv2 +from PIL import Image +import os +import json +from shapely.geometry import Polygon, MultiPolygon +from skimage.draw import polygon +from skimage.morphology import erosion, square, binary_erosion +from skimage.io import imread,imsave +from utils import colorize,create_path +from tqdm import tqdm + +import rasterio as rs +from cv2 import fillPoly,copyMakeBorder +from math import ceil + + +class Masker(): + def __init__(self, + data, + out_size = (1024,1024), + erosion_kernel ='cross', + iterator_verbose = True): + self.sz = out_size + self.pd_sz = (1044,1044) + self.x_off = ceil((self.pd_sz[1] - self.sz[1])/2) + self.y_off = ceil((self.pd_sz[0] - self.sz[0])/2) + assert self.x_off >= 0 and self.y_off >= 0, f'out size {self.sz} should be less than padded size {self.pd_sz}' + assert erosion_kernel.lower() in ['square','cross'], f'erosion kernel type : [ {erosion_kernel} ] is not valid' + self.ek_type = erosion_kernel.lower() + self.itr_vrbs = iterator_verbose ### + + self.data = data + self.ids = sorted(os.listdir(self.data)) + self.ldir = 'labels_match_pix' + + + def load_labels(self,json_path): + jfile=open(json_path,'r') + f=json.load(jfile) + return f + + def poly_size(self,w,h): + return Polygon([[0,0],[w-1,0],[w-1,h-1],[0,h-1],[0,0]]) + + def get_strc(self): + if(self.ek_type == 'square'): + return square(3) + else: + return np.array([[0,1,0], + [1,1,1], + [0,1,0]],dtype=np.uint8) + + + + + def load_raster_file(self,raster_path): + return rs.open(raster_path) + + def get_img(self,raster_file): + return raster_file.read().transpose(1,2,0).astype(np.uint8) + + def project_poly(self,poly,frs,size,x_off= 0,y_off=0): + k = [] + for tup in poly: + _x,_y = frs.index(*tup) + _x+=x_off + _y+=y_off + k.append([_x,_y]) + + polk= Polygon(k) + if(not polk.is_valid): + polk = polk.buffer(0.01) + poll = self.poly_size(*size) + intr = poll.intersection(polk) + verbs = intr.geom_type == 'Polygon' + return list(intr.exterior.coords) if(verbs) else [] + + def crop(self,img,y_off,x_off,h,w): + return img[y_off:y_off + h,x_off:x_off + w] + + def make_mask(self,polys, size = (1024,1024)): + builds,border = [np.zeros(size,dtype=np.uint8),np.zeros(size,dtype=np.uint8)] + instances = np.zeros(size,dtype=np.int32) + strc = self.get_strc() + itr = enumerate(polys) + if(self.itr_vrbs): + itr = tqdm(itr) + itr.set_description('generating mask') + for i,mulpol in itr: + for j,pol in enumerate(mulpol): + arr_pol= np.array(pol,dtype=np.int32) + hs,ws = polygon(arr_pol[:,0],arr_pol[:,1],size) + instances[hs,ws,...] = np.int32(i+1) if(j==0) else 0 + #fillPoly(instances, [arr_pol[:,::-1]], i+1 if(j == 0) else 0) + instance = instances == np.int32(i+1) + try: + k=np.where(instance>0) + _t = k[0].min() - 2 + _l = k[1].min() - 2 + _b = k[0].max() + 2 + _r = k[1].max() + 2 + + crop_instance = instance[_t:_b,_l:_r] + bld = binary_erosion(crop_instance,selem = strc) + brdr = bld ^ crop_instance + brdr1 = np.zeros_like(instance,dtype=brdr.dtype) + brdr1[_t:_b,_l:_r] =brdr + border[brdr1 == True] = np.uint8(255) + + except: + bld = binary_erosion(instance,selem = strc) + brdr = bld ^ instance + border[brdr == True] = np.uint8(255) + + builds[instances > 0] = np.uint8(255) + return instances,builds,border + + def mask(self,raster_path,json_path): + raster = self.load_raster_file(raster_path) + img = self.get_img(raster) + js = self.load_labels(json_path) + labels = js['features'] + polys =[] + for label in labels: + + multipoly = label['geometry']['coordinates'] + proj_multipoly = [] + for poly in multipoly: + mm = self.project_poly(poly, raster,self.pd_sz,self.x_off,self.y_off) + if(len(mm)>0): + proj_multipoly.append(mm) + polys.append(proj_multipoly) + + ins,b,br = self.make_mask(polys,size = self.pd_sz) + kwargs = {'y_off' : self.y_off, + 'x_off' : self.x_off, + 'h' : self.sz[0], + 'w' : self.sz[1]} + ins = self.crop(ins,**kwargs) + b = self.crop(b,**kwargs) + br = self.crop(br,**kwargs) + img = cv2.copyMakeBorder(img, + top=0, + bottom = self.sz[0] - img.shape[0], + left=0, + right = self.sz[1] - img.shape[1], + borderType =cv2.BORDER_CONSTANT, + value = 0) + return img,ins,b,br + + + + + def _collect(self,labels): + _meta = {} + for label in labels['features']: + pid = str(np.int32(label['properties']['Id'])) + _meta[pid] = {} + _meta[pid]['area']= label['properties']['area'] + _meta[pid]['geometry'] = label['geometry']['coordinates'] + + def int_coords(self,x): + return np.array(x).round().astype(np.int32) + + def instances(self,size,labels): + ins_mask = np.zeros(size,dtype = np.int32) + + for pid,d in labels.items(): + int_id = np.int32(pid) + polys = d['geometry'] + for i,poly in enumerate(polys): + poly.append(poly[0]) + S = Polygon(poly) + PS = self.poly_size(size[1],size[0]) + S = S.intersection(PS) + Stype = S.geom_type + + if(Stype == 'Polygon'): + arr_pol = self.int_coords(S.exterior.coords) + if(len(arr_pol.shape) != 2): continue + ws,hs = polygon(arr_pol[:,0],arr_pol[:,1],size) + ins_mask[hs,ws,...] = int_id if(i==0) else 0 + + elif(Stype == 'MultiPolygon'): + for s in S: + arr_pol = self.int_coords(s.exterior.coords) + if(len(arr_pol.shape) != 2): continue + ws,hs = polygon(arr_pol[:,0],arr_pol[:,1],size) + ins_mask[hs,ws,...] = int_id if(i==0) else 0 + + else: + for point in list(S.coords): + x,y = list(map(np.int32,point)) + ins_mask[y,x,...] = int_id if(i==0) else 0 + + return ins_mask + + def borders(self,ins_mask): + ins_borders = np.zeros_like(ins_mask,dtype = np.int32) + ids = sorted(np.unique(ins_mask))[1:] + strc = self.get_strc() + for iid in ids: + instance = ins_mask == iid + try: + k=np.where(instance>0) + _t = k[0].min() - 3 + _l = k[1].min() - 3 + _b = k[0].max() + 3 + _r = k[1].max() + 3 + + crop_instance = instance[_t:_b,_l:_r] + bld = binary_erosion(crop_instance,selem = strc) + brdr = bld ^ crop_instance + brdr1 = np.zeros_like(instance,dtype=brdr.dtype) + brdr1[_t:_b,_l:_r] =brdr + ins_borders[brdr1 == True] = iid + + except: + bld = binary_erosion(instance,selem = strc) + brdr = bld ^ instance + ins_borders[brdr == True] = iid + return ins_borders + + def to_rgb(self,img): + rgb = np.ascontiguousarray(img[...,:3],dtype = np.uint8) + return rgb + + def to_gray(self,mask): + return (mask>0).astype(np.uint8) * 255 + + def generate_dataset(self,save_path): + create_path(save_path) + + for iid in self.ids[42:]: + imgs_save = f'{save_path}/{iid}' + imgs_path = f'{self.data}/{iid}/images_masked' + labels_path = f'{self.data}/{iid}/{self.ldir}' + lod = os.listdir(imgs_path) + loader = tqdm(lod) if(itr_vrbs) else lod + loader.set_description(f'{iid}') + create_path(imgs_save) + for exten in loader: + e = exten.split('.tif')[0] + img_save = f'{imgs_save}/{e}' + create_path(img_save) + + img_rgba = imread(f'{imgs_path}/{e}.tif') + img_rgb = self.to_rgb(img_rgba) + + lf = self.load_labels(iid, e) + labels = self._collect(lf) + + size = img_rgb.shape[:2] + ins_mask = self.instances(size, labels) + ins_borders = self.borders(ins_mask) + + imsave(f'{img_save}/image.png',img_rgb) + imsave(f'{img_save}/buildings.png',self.to_gray(ins_mask)) + imsave(f'{img_save}/borders.png',self.to_gray(ins_borders)) + np.save(f'{img_save}/buildings.npy',ins_mask) + np.save(f'{img_save}/borders.npy',ins_borders) + + +# class masker(): +# def __init__(self,data, +# out_size = (1024,1024), +# erosion_kernel ='cross', +# verbose = True): +# self.data = data +# self.sz = out_size +# self.ids = sorted(os.listdir(self.data)) +# self.ldir = 'labels_match_pix' +# assert erosion_kernel in ['cross','square'] +# self.ek_type = erosion_kernel +# itr_vrbs = verbose + \ No newline at end of file