学习小记：语义分割模型训练后的测试及miou值的求取（python）

小白学习语义分割，记录一下，方便以后查阅。
借鉴了这个链接里面的deepla b-mobilenet模型：
https://github.com/bubbliiiing/Semantic-Segmentation
求miou主要借鉴这个链接：
https://blog.csdn.net/weixin_42750737/article/details/88886242
将miou的求解结合到了第一个链接里面的predict中，用于评估语义分割的模型。

from nets.deeplab import Deeplabv3
from PIL import Image
import numpy as np
import random
import copy
import os
import time
s=time.time()
#标签颜色
class_colors = [[0,0,0],[255, 0, 0],[255, 255, 0],
        [0, 255, 0],
        [0,255,255],
        [148,0,211],
        [255,140,0],
        [83,83,0]]
NCLASSES = 8
HEIGHT = 513
WIDTH = 513
model =  Deeplabv3(classes=NCLASSES,input_shape=(HEIGHT,WIDTH,3))
model.load_weights("./dataset2/last1.h5")
imgs = os.listdir("./dataset2/test_jpg") #原图文件
labels="./dataset2/test_png/"  #标签文件
txt=open("./miou.txt",'r+')
def fast_hist(label,prediction,n):
    #去掉忽略类，有的说去掉了背景，我认为没有去掉，背景设置的类为0
    k = (label >= 0) & (label < n)
    # k为掩码，保留a中对应k中为true的值，[0,n-1]对应标签的第一类被预测成的类，[n,2n-1]对应标签第二类被预测成的类，以此类推。
    return np.bincount(n * label[k] + prediction[k], minlength=n** 2).reshape(n, n)
def per_class_iu(hist):
    #iou=TP/(TP+FN+FP)，对角线表示的是TP，sum(1)每行相加，表示实际为某类的所有像素，sum(0)每列相加，表示预测为某类的所有像素。
    #横向为预测，纵向为实际，返回的为一个向量，每一类的miou
    iou=np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
    return iou
def compute_miou(label,prediction,n):
    hist=fast_hist(label,prediction,n)
    iou=per_class_iu(hist)
    return iou
hist=np.zeros((NCLASSES, NCLASSES))
for jpg in imgs:
    img = Image.open("./dataset2/test_jpg/"+jpg)
    old_img = copy.deepcopy(img)
    orininal_h = np.array(img).shape[0]
    orininal_w = np.array(img).shape[1]
    img = img.resize((WIDTH,HEIGHT))  #x,y方向
    img = np.array(img)
    img = img/255
    img = img.reshape(-1,HEIGHT,WIDTH,3)
    pr = model.predict(img)[0]  #输出为类别概率图
    #判断每一个像素属于哪一类，在哪个类别的概率最大，.argmax最大值的索引
    pr = pr.reshape((int(HEIGHT), int(WIDTH),NCLASSES)).argmax(axis=-1)
    #计算IOU
    prediction=pr.reshape((-1))
    label1=Image.open(labels+jpg).convert('L')
    #选择最近邻插值缩放，双线性插值缩放会改变标签图中像素值
    label1=cv2.resize(label1, (HEIGHT,WIDTH), interpolation=cv2.INTER_NEAREST)
    label1=np.array(label1)
    #print(np.unique(label1))
    label=label1.reshape((-1))
    hist+=fast_hist(label,prediction,NCLASSES)
    iou=compute_miou(label,prediction,NCLASSES) #单张图片的各类iou
    miou=np.nanmean(iou)
    print(iou)
#将结果写入txt文件中分析
    txt.truncate()
    txt.write("---------*"+jpg+"*---------""\n")
    txt.write("The real classes are \n")
    txt.write("\t"+str(np.unique(label1))+"\n")
    txt.write("The predicted classes are \n" )
    txt.write("\t"+str(np.unique(pr))+"\n")
    txt.write("The iou value for each category in single picture is \n")
    txt.write("\t"+str([round(x*100,2) for x in iou])+"\n")#round（x，n） 小数点后n位四舍五入
    txt.write("The single picture miou is"+"\t"+str(round(miou*100,2))+"\n"+"\n")
    #对每个类别着色
    seg_img = np.zeros((int(HEIGHT), int(WIDTH),3))
    colors = class_colors
    for c in range(NCLASSES):
        seg_img[:,:,0] += ( (pr[:,: ] == c )*( colors[c][0] )).astype('uint8')
        seg_img[:,:,1] += ((pr[:,: ] == c )*( colors[c][1] )).astype('uint8')
        seg_img[:,:,2] += ((pr[:,: ] == c )*( colors[c][2] )).astype('uint8')
    #array转化为image
    seg_img = Image.fromarray(np.uint8(seg_img)).resize((orininal_w,orininal_h))
    seg_img.save("./img_out/1" + jpg)
    image = Image.blend(old_img,seg_img,0.3)
    #blended_img = img1 * (1 – alpha) + img2* alpha
    image.save("./img_out/2"+jpg)
#测试集总的miou
all_iou=per_class_iu(hist)
all_miou=np.nanmean(all_iou)*100
n_img=len(imgs)
e=time.time()
t=e-s
print("miou:%.2f"%all_miou)
#在txt文件头部写入总的miou
txt.close()
with open("./miou.txt",'r+') as f:
    content=f.read()
    f.seek(0,0)
    head="It takes %f seconds to predict %d images"%(t,n_img)+"\n"\
          "The miou in test dataset is :%.2f " % all_miou+"\n"\
         "The iou value for each category in test dataset are "+str([round(y*100,2) for y in all_iou])+"\n"
    f.write(head+content)
print("It takes %f seconds to predict %d images"%(t,n_img))