import torch from core.models.bisenet import BiSeNet,BiSeNet_MultiOutput from torchvision import transforms import cv2,os,glob import numpy as np from core.models.dinknet import DinkNet34 import matplotlib.pyplot as plt import matplotlib.pyplot as plt import time class SegModel(object): def __init__(self, nclass=2,model = None,weights=None,modelsize=512,device='cuda:3',multiOutput=False): #self.args = args self.model = model #self.model = DinkNet34(nclass) checkpoint = torch.load(weights) self.modelsize = modelsize self.model.load_state_dict(checkpoint['model']) self.device = device self.multiOutput = multiOutput self.model= self.model.to(self.device) '''self.composed_transforms = transforms.Compose([ transforms.Normalize(mean=(0.335, 0.358, 0.332), std=(0.141, 0.138, 0.143)), transforms.ToTensor()]) ''' self.mean = (0.335, 0.358, 0.332) self.std = (0.141, 0.138, 0.143) #mean=(0.335, 0.358, 0.332), std=(0.141, 0.138, 0.143) def eval(self,image,outsize=None,smooth_kernel=0): imageH,imageW,imageC = image.shape time0 = time.time() image = self.preprocess_image(image) time1 = time.time() self.model.eval() image = image.to(self.device) with torch.no_grad(): output = self.model(image,test_flag=True,smooth_kernel = 0) time2 = time.time() if self.multiOutput: pred = [outputx.data.cpu().numpy()[0] for outputx in output] else: pred = output.data.cpu().numpy() pred = pred[0] time3 = time.time() if self.multiOutput: pred = [ cv2.blur(predx,(smooth_kernel,smooth_kernel) ) for predx in pred] pred = [cv2.resize(predx.astype(np.uint8),(imageW,imageH)) for predx in pred[0:2]] else: pred = cv2.blur(pred,(smooth_kernel,smooth_kernel) ) pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH),interpolation = cv2.INTER_NEAREST) time4 = time.time() print('##line52:pre-precess:%.1f ,infer:%.1f ,post-precess:%.1f ,post-resize:%.1f '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3) )) return pred def get_ms(self,t1,t0): return (t1-t0)*1000.0 def preprocess_image(self,image): time0 = time.time() image = cv2.resize(image,(self.modelsize,self.modelsize)) time1 = time.time() image = image.astype(np.float32) image /= 255.0 time2 = time.time() #image = image * 3.2 - 1.6 image[:,:,0] -=self.mean[0] image[:,:,1] -=self.mean[1] image[:,:,2] -=self.mean[2] time3 = time.time() image[:,:,0] /= self.std[0] image[:,:,1] /= self.std[1] image[:,:,2] /= self.std[2] time4 = time.time() image = np.transpose(image, ( 2, 0, 1)) time5 = time.time() image = torch.from_numpy(image).float() image = image.unsqueeze(0) print('###line84: in preprocess: resize:%.1f norm:%.1f mean:%.1f std:%.1f trans:%.f '%(self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3) ,self.get_ms(time5,time4) ) ) return image def get_ms(t1,t0): return (t1-t0)*1000.0 def test(): #os.environ["CUDA_VISIBLE_DEVICES"] = str('4') ''' image_url = '../../data/landcover/corp512/test/images/N-33-139-C-d-2-4_169.jpg' nclass = 5 weights = 'runs/landcover/DinkNet34_save/experiment_wj_loss-10-10-1/checkpoint.pth' ''' image_url = 'temp_pics/DJI_0645.JPG' nclass = 2 #weights = '../weights/segmentation/BiSeNet/checkpoint.pth' weights = 'runs/THriver/BiSeNet/train/experiment_0/checkpoint.pth' #weights = 'runs/segmentation/BiSeNet_test/experiment_10/checkpoint.pth' model = BiSeNet(nclass) segmodel = SegModel(model=model,nclass=nclass,weights=weights,device='cuda:4') for i in range(10): image_array0 = cv2.imread(image_url) imageH,imageW,_ = image_array0.shape #print('###line84:',image_array0.shape) image_array = cv2.cvtColor( image_array0,cv2.COLOR_RGB2BGR) #image_in = segmodel.preprocess_image(image_array) pred = segmodel.eval(image_array,outsize=None) time0=time.time() binary = pred.copy() time1=time.time() contours, hierarchy = cv2.findContours(binary,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) time2=time.time() print(pred.shape,' time copy:%.1f finccontour:%.1f '%(get_ms(time1,time0),get_ms(time2,time1) )) label_dic={'landcover':[[0, 0, 0], [255, 0, 0], [0,255,0], [0,0,255], [255,255,0]], 'deepRoad':[[0,0,0],[255,0,0]], 'water':[[0,0,0],[255,255,255]], 'water_building':[[0,0,0],[0,0,255],[255,0,0]], 'floater':[[0,0,0], [0,255,0],[255,255,0],[255,0,255],[0,128, 255], [255,0,0], [0,255,255] ] } def index2color(label_mask,label_colours): r = label_mask.copy() g = label_mask.copy() b = label_mask.copy() label_cnt = len(label_colours) for ll in range(0, label_cnt): r[label_mask == ll] = label_colours[ll][0] g[label_mask == ll] = label_colours[ll][1] b[label_mask == ll] = label_colours[ll][2] rgb = np.stack((b, g,r), axis=-1) return rgb.astype(np.uint8) def get_largest_contours(contours): areas = [cv2.contourArea(x) for x in contours] max_area = max(areas) max_id = areas.index(max_area) return max_id def result_merge_sep(image,mask_colors): #mask_colors=[{ 'mask':mask_map,'index':[1],'color':[255,255,255] }] for mask_color in mask_colors: mask_map,indexes,colors = mask_color['mask'], mask_color['index'], mask_color['color'] ishow = 2 #plt.figure(1);plt.imshow(mask_map); for index,color in zip(indexes,colors): mask_binaray = (mask_map == index).astype(np.uint8) contours, hierarchy = cv2.findContours(mask_binaray,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) if len(contours)>0: d=hierarchy[0,:,3]<0 ; contours = np.array(contours,dtype=object)[d] cv2.drawContours(image,contours,-1,color[::-1],3) #plt.figure(ishow);plt.imshow(mask_binaray);ishow+=1 #plt.show() return image def result_merge(image,mask_colors): #mask_colors=[{ 'mask':mask_map,'index':[1],'color':[255,255,255] }] for mask_color in mask_colors: mask_map,indexes,colors = mask_color['mask'], mask_color['index'], mask_color['color'] mask_binary = (mask_map>0).astype(np.uint8) contours, hierarchy = cv2.findContours(mask_binary,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) if len(contours)>0: d=hierarchy[0,:,3]<0 ; contours = np.array(contours)[d] cv2.drawContours(image,contours,-1,colors[0][::-1],3) coors = np.array([(np.mean(contours_x ,axis=0)+0.5).astype(np.int32)[0] for contours_x in contours]) #print(mask_map.shape,coors.shape) typess = mask_map[ coors[:,1],coors[:,0]] #for jj,iclass in enumerate(typess): #print(iclass,colors) # cv2.drawContours(image,contours,-1, colors[iclass][::-1],3) return image def test_floater(): from core.models.dinknet import DinkNet34_MultiOutput #create_model('DinkNet34_MultiOutput',[2,5]) image_url = 'temp_pics/DJI_0645.JPG' nclass = [2,7] outresult=True weights = 'runs/thFloater/BiSeNet_MultiOutput/train/experiment_4/checkpoint.pth' model = BiSeNet_MultiOutput(nclass) outdir='temp' image_dir = '/host/workspace/WJ/data/thFloater/val/images/' image_url_list=glob.glob('%s/*'%(image_dir)) segmodel = SegModel(model=model,nclass=nclass,weights=weights,device='cuda:9',multiOutput=True) for i,image_url in enumerate(image_url_list[0:10]) : image_array0 = cv2.imread(image_url) image_array0 = cv2.cvtColor(image_array0, cv2.COLOR_BGR2RGB) # cv2默认为bgr顺序 imageH,imageW,_ = image_array0.shape #image_array = cv2.cvtColor( image_array0,cv2.COLOR_RGB2BGR) pred = segmodel.eval(image_array,outsize=None) time0=time.time() if isinstance(pred,list): binary = [predx.copy() for predx in pred] time1=time.time() mask_colors=[ { 'mask':pred[0] ,'index':range(1,2),'color':label_dic['water'][0:] }, { 'mask':pred[1] ,'index':[1,2,3,4,5,6],'color':label_dic['floater'][0:] } ] result_draw = result_merge(image_array0,mask_colors) time2=time.time() if outresult: basename=os.path.splitext( os.path.basename(image_url))[0] outname=os.path.join(outdir,basename+'_draw.png') cv2.imwrite(outname,result_draw[:,:,:]) print('##line151: time copy:%.1f finccontour:%.1f '%(get_ms(time1,time0),get_ms(time2,time1) )) def test_water_buildings(): from core.models.bisenet import BiSeNet #image_url = 'temp_pics/DJI_0645.JPG' nclass = 3 outresult=True weights = 'runs/thWaterBuilding/BiSeNet/train/experiment_2/checkpoint.pth' model = BiSeNet(nclass) outdir='temp' image_dir = '/home/thsw/WJ/data/river_buildings/' #image_dir = '/home/thsw/WJ/data/THWaterBuilding/val/images' image_url_list=glob.glob('%s/*'%(image_dir)) segmodel = SegModel(model=model,nclass=nclass,weights=weights,device='cuda:0',multiOutput=False) for i,image_url in enumerate(image_url_list[0:]) : #image_url = '/home/thsw/WJ/data/THWaterBuilding/val/images/0anWqgmO9rGe1n8P.png' image_array0 = cv2.imread(image_url) imageH,imageW,_ = image_array0.shape image_array = cv2.cvtColor( image_array0,cv2.COLOR_RGB2BGR) pred = segmodel.eval(image_array,outsize=None) time0=time.time() if isinstance(pred,list): binary = [predx.copy() for predx in pred] #print(binary[0].shape) time1=time.time() mask_colors=[ { 'mask':pred ,'index':range(1,3),'color':label_dic['water_building'][1:] }, #{ 'mask':pred[1] ,'index':[1,2,3,4,5,6],'color':label_dic['floater'][0:] } ] result_draw = result_merge_sep(image_array0,mask_colors) time2=time.time() if outresult: basename=os.path.splitext( os.path.basename(image_url))[0] outname=os.path.join(outdir,basename+'_draw.png') cv2.imwrite(outname,result_draw[:,:,:]) print('##line294: time copy:%.1f finccontour:%.1f '%(get_ms(time1,time0),get_ms(time2,time1) )) def get_illegal_index(contours,hierarchy,water_dilate,overlap_threshold): out_index=[] if len(contours)>0: d=hierarchy[0,:,3]<0 ; contours = np.array(contours,dtype=object)[d] imageH,imageW = water_dilate.shape for ii,cont in enumerate(contours): build_area=np.zeros((imageH,imageW )) cv2.fillPoly(build_area,[cont[:,0,:]],1) area1=np.sum(build_area);area2=np.sum(build_area*water_dilate) if (area2/area1) >overlap_threshold: out_index.append(ii) return out_index def test_water_building_seperately(): from core.models.dinknet import DinkNet34_MultiOutput #create_model('DinkNet34_MultiOutput',[2,5]) image_url = 'temp_pics/DJI_0645.JPG' nclass = [2,2] outresult=True weights = 'runs/thWaterBuilding_seperate/BiSeNet_MultiOutput/train/experiment_0/checkpoint.pth' model = BiSeNet_MultiOutput(nclass) outdir='temp' image_dir = '/home/thsw/WJ/data/river_buildings/' #image_dir = '/home/thsw/WJ/data/THWaterBuilding/val/images' image_url_list=glob.glob('%s/*'%(image_dir)) segmodel = SegModel(model=model,nclass=nclass,weights=weights,device='cuda:1',multiOutput=True) print('###line307 image cnt:',len(image_url_list)) for i,image_url in enumerate(image_url_list[0:1]) : image_url = '/home/thsw/WJ/data/river_buildings/DJI_20210904092044_0001_S_output896.jpg' image_array0 = cv2.imread(image_url) imageH,imageW,_ = image_array0.shape image_array = cv2.cvtColor( image_array0,cv2.COLOR_RGB2BGR) pred = segmodel.eval(image_array,outsize=None,smooth_kernel=20) ##画出水体区域 contours, hierarchy = cv2.findContours(pred[0],cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) max_id = get_largest_contours(contours); water = pred[0].copy(); water[:,:] = 0 cv2.fillPoly(water, [contours[max_id][:,0,:]], 1) cv2.drawContours(image_array0,contours,max_id,(0,255,255),3) ##画出水体膨胀后的蓝线区域。 kernel = np.ones((100,100),np.uint8) water_dilate = cv2.dilate(water,kernel,iterations = 1) contours, hierarchy = cv2.findContours(water_dilate,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) #print('####line310:',contours) cv2.drawContours(image_array0,contours,-1,(255,0,0),3) ###逐个建筑判断是否与蓝线内区域有交叉。如果交叉面积占本身面积超过0.1,则认为是违法建筑。 contours, hierarchy = cv2.findContours(pred[1],cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) outIndex=get_illegal_index(contours,hierarchy,water_dilate,0.1) for ii in outIndex: cv2.drawContours(image_array0,contours,ii,(0,0,255),3) plt.imshow(image_array0);plt.show() ## time0=time.time() time1=time.time() mask_colors=[ { 'mask':pred[0],'index':[1],'color':label_dic['water_building'][1:2]}, { 'mask':pred[1],'index':[1],'color':label_dic['water_building'][2:3]} ] result_draw = result_merge_sep(image_array0,mask_colors) time2=time.time() if outresult: basename=os.path.splitext( os.path.basename(image_url))[0] outname=os.path.join(outdir,basename+'_draw.png') cv2.imwrite(outname,result_draw[:,:,:]) print('##line151: time copy:%.1f finccontour:%.1f '%(get_ms(time1,time0),get_ms(time2,time1) )) if __name__=='__main__': #test() #test_floater() #test_water_buildings() test_water_building_seperately()