132 lines
4.3 KiB
Python
132 lines
4.3 KiB
Python
import torch
|
|
import sys,os
|
|
sys.path.extend(['../AIlib/segutils'])
|
|
from model_stages import BiSeNet_STDC
|
|
from torchvision import transforms
|
|
import cv2,glob
|
|
import numpy as np
|
|
from core.models.dinknet import DinkNet34
|
|
import matplotlib.pyplot as plt
|
|
import time
|
|
class SegModel(object):
|
|
def __init__(self, nclass=2,weights=None,modelsize=(640,360),device='cuda:0'):
|
|
#self.args = args
|
|
self.model = BiSeNet_STDC(backbone='STDCNet813', n_classes=nclass,
|
|
use_boundary_2=False, use_boundary_4=False,
|
|
use_boundary_8=True, use_boundary_16=False,
|
|
use_conv_last=False)
|
|
self.device = device
|
|
self.model.load_state_dict(torch.load(weights, map_location=torch.device(self.device) ))
|
|
self.model= self.model.to(self.device)
|
|
self.mean = (0.485, 0.456, 0.406)
|
|
self.std = (0.229, 0.224, 0.225)
|
|
self.modelsize=modelsize
|
|
|
|
def eval(self,image):
|
|
time0 = time.time()
|
|
imageH, imageW, _ = image.shape
|
|
image = self.RB_convert(image)
|
|
img = self.preprocess_image(image)
|
|
if self.device != 'cpu':
|
|
imgs = img.to(self.device)
|
|
else:imgs=img
|
|
time1 = time.time()
|
|
self.model.eval()
|
|
with torch.no_grad():
|
|
output = self.model(imgs)
|
|
|
|
time2 = time.time()
|
|
pred = output.data.cpu().numpy()
|
|
pred = np.argmax(pred, axis=1)[0]#得到每行
|
|
time3 = time.time()
|
|
pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH))
|
|
time4 = time.time()
|
|
outstr= 'pre-precess:%.1f ,infer:%.1f ,post-cpu-argmax:%.1f ,post-resize:%.1f, total:%.1f \n '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
|
|
|
|
return pred,outstr
|
|
def get_ms(self,t1,t0):
|
|
return (t1-t0)*1000.0
|
|
def preprocess_image(self,image):
|
|
image = cv2.resize(image, self.modelsize, interpolation=cv2.INTER_LINEAR)
|
|
image = image.astype(np.float32)
|
|
image /= 255.0
|
|
|
|
image[:, :, 0] -= self.mean[0]
|
|
image[:, :, 1] -= self.mean[1]
|
|
image[:, :, 2] -= self.mean[2]
|
|
|
|
image[:, :, 0] /= self.std[0]
|
|
image[:, :, 1] /= self.std[1]
|
|
image[:, :, 2] /= self.std[2]
|
|
|
|
image = np.transpose(image, (2, 0, 1))
|
|
image = torch.from_numpy(image).float()
|
|
image = image.unsqueeze(0)
|
|
|
|
return image
|
|
def RB_convert(self,image):
|
|
image_c = image.copy()
|
|
image_c[:,:,0] = image[:,:,2]
|
|
image_c[:,:,2] = image[:,:,0]
|
|
return image_c
|
|
|
|
def get_ms(t1,t0):
|
|
return (t1-t0)*1000.0
|
|
|
|
|
|
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
|
|
|
|
if __name__=='__main__':
|
|
impth = '../../../../data/无人机起飞测试图像/'
|
|
outpth= 'results'
|
|
folders = os.listdir(impth)
|
|
weights = '../weights/STDC/model_maxmIOU75_1720_0.946_360640.pth'
|
|
segmodel = SegModel(nclass=2,weights=weights)
|
|
|
|
for i in range(len(folders)):
|
|
|
|
imgpath = os.path.join(impth, folders[i])
|
|
time0 = time.time()
|
|
|
|
#img = Image.open(imgpath).convert('RGB')
|
|
img = cv2.imread(imgpath)
|
|
img = np.array(img)
|
|
|
|
|
|
time1 = time.time()
|
|
pred, outstr = segmodel.eval(image=img)#####
|
|
time2 = time.time()
|
|
|
|
binary0 = pred.copy()
|
|
contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
|
|
time3 = time.time()
|
|
|
|
max_id = -1
|
|
if len(contours)>0:
|
|
max_id = get_largest_contours(contours)
|
|
binary0[:,:] = 0
|
|
cv2.fillPoly(binary0, [contours[max_id][:,0,:]], 1)
|
|
|
|
cv2.drawContours(img,contours,max_id,(0,255,255),3)
|
|
time4 = time.time()
|
|
|
|
|
|
#img_n = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
|
|
cv2.imwrite( os.path.join( outpth,folders[i] ) ,img )
|
|
time5 = time.time()
|
|
print('image:%d ,infer:%.1f ms,findcontours:%.1f ms, draw:%.1f, total:%.1f'%(i,get_ms(time2,time1),get_ms(time3,time2),get_ms(time4,time3),get_ms(time4,time1)))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|