Pytorch 加載保存模型，進行模型推斷【直播】2019 年縣域農業大腦AI挑戰賽---(三)保存結果

在模型訓練結束，結束后，通常是一個分割模型，輸入 1024x1024 輸出 4x1024x1024。

一種方法就是將整個圖切塊，然后每張預測，但是有個不好處就是可能在邊界處斷續。

 

由於這種切塊再預測很ugly，所以直接遍歷整個圖預測(這就是相當於卷積啊)，防止邊界斷續，還有一個問題就是防止圖過大不能超過20M。

很有意思解決上邊的問題。話也不多說了。直接上代碼：

from farmlanddataset import FarmDataset
import torch as tc
from osgeo import gdal
from torchvision import transforms
import png
import numpy as np
use_cuda=True
model=tc.load('./tmp/model30')  #torch.save(model,'./tmp/model{}'.format(epoch))
device = tc.device("cuda" if use_cuda else "cpu")
model=model.to(device)
model.eval()
ds=FarmDataset(istrain=False)

def createres(d,outputname):
    #創建一個和ds大小相同的灰度圖像BMP
    driver = gdal.GetDriverByName("BMP")
    #driver=ds.GetDriver()
    od=driver.Create('./tmp/'+outputname,d.RasterXSize,d.RasterYSize,1)
    return od

def createpng(height,width,data,outputname):
    w=png.Writer(width,height,bitdepth=2,greyscale=True)
    of=open('./tmp/'+outputname,'wb')
    w.write_array(of,data.flat)
    of.close()
    return 
def predict(d,outputname='tmp.bmp'):
    wx=d.RasterXSize   #width
    wy=d.RasterYSize   #height
    print(wx,wy)
    od=data=np.zeros((wy,wx),np.uint8)
    #od=createres(d,outputname=outputname)
    #ob=od.GetRasterBand(1) #得到第一個channnel
    blocksize=1024
    step=512
    for cy in range(step,wy-blocksize,step):
        for cx in range(step,wx-blocksize,step):
            img=d.ReadAsArray(cx-step,cy-step,blocksize,blocksize)[0:3,:,:] #channel*h*w
            if (img.sum()==0): continue
            x=tc.from_numpy(img/255.0).float()
            #print(x.shape)
            x=x.unsqueeze(0).to(device)
            r=model.forward(x)
            r=tc.argmax(r.cpu()[0],0).byte().numpy()  #512*512
            #ob.WriteArray(r,cx,cy)
            od[cy-step//2:cy+step//2,cx-step//2:cx+step//2]=r[256:step+256,256:step+256]
            print(cy,cx)
    #del od
    createpng(wy,wx,od,outputname)
    return 
    
print("start predict.....")
predict(ds[0],'image_3_predict.png')
print("start predict 2 .....")
predict(ds[1],'image_4_predict.png')

　　

然后看看我的結果：提交了，晚上希望有個不錯的結果

看上邊的分類結果，真是感慨深度學習大法好，傳統的遙感分類完全沒有辦法，上邊結果在比賽中評測指標>0.2。

 有了這些就可以發揮想象力和搬家能力，訓練模型。