108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
-------------------------------------------------
|
|||
|
|
Project Name: unet
|
|||
|
|
File Name: test.py
|
|||
|
|
Author: chenming
|
|||
|
|
Create Date: 2022/2/7
|
|||
|
|
Description：
|
|||
|
|
-------------------------------------------------
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
from utils.utils_metrics import compute_mIoU, show_results
|
|||
|
|
import glob
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
import os
|
|||
|
|
import cv2
|
|||
|
|
from model.unet_model import UNet
|
|||
|
|
from torchmetrics import Dice
|
|||
|
|
|
|||
|
|
def cal_miou(test_dir=r"C:\Users\USER\Desktop\segment\test\Testimages_new",
|
|||
|
|
pred_dir=r"C:\Users\USER\Desktop\segment\unet_42-master\results",
|
|||
|
|
gt_dir=r"C:\Users\USER\Desktop\segment\test\Testlabels_new"):
|
|||
|
|
# ---------------------------------------------------------------------------#
|
|||
|
|
# miou_mode 用於指定該文件運行時計算的內容
|
|||
|
|
# miou_mode 為 0 代表整個 miou 計算流程，包括獲得預測結果、計算 miou。
|
|||
|
|
# miou_mode 為 1 代表僅僅獲得預測結果。
|
|||
|
|
# miou_mode 為 2 代表僅僅計算 miou。
|
|||
|
|
# ---------------------------------------------------------------------------#
|
|||
|
|
miou_mode = 0
|
|||
|
|
# ------------------------------#
|
|||
|
|
# 分類個數 + 1，如 2 + 1
|
|||
|
|
# ------------------------------#
|
|||
|
|
num_classes = 2
|
|||
|
|
# --------------------------------------------#
|
|||
|
|
# 區分的種類，和 json_to_dataset 裡面的一樣
|
|||
|
|
# --------------------------------------------#
|
|||
|
|
name_classes = ["background", "hippocampus"]
|
|||
|
|
# name_classes = ["_background_","cat","dog"]
|
|||
|
|
# -------------------------------------------------------#
|
|||
|
|
# 指向 VOC 數據集所在的文件夾
|
|||
|
|
# 預設指向根目錄下的 VOC 數據集
|
|||
|
|
# -------------------------------------------------------#
|
|||
|
|
# 計算結果與 gt 的結果進行比對
|
|||
|
|
|
|||
|
|
# 載入模型
|
|||
|
|
|
|||
|
|
if miou_mode == 0 or miou_mode == 1:
|
|||
|
|
if not os.path.exists(pred_dir):
|
|||
|
|
os.makedirs(pred_dir)
|
|||
|
|
|
|||
|
|
print("載入模型中。")
|
|||
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|||
|
|
# 載入網路，圖片為單通道，分類為 1。
|
|||
|
|
net = UNet(n_channels=1, n_classes=1)
|
|||
|
|
# 將網路拷貝至裝置中
|
|||
|
|
net.to(device=device)
|
|||
|
|
# 載入模型參數
|
|||
|
|
net.load_state_dict(torch.load('best_model_30_1.pth', map_location=device)) # todo
|
|||
|
|
# 測試模式
|
|||
|
|
net.eval()
|
|||
|
|
print("模型載入完成。")
|
|||
|
|
|
|||
|
|
img_names = os.listdir(test_dir)
|
|||
|
|
image_ids = [image_name.split(".")[0] for image_name in img_names]
|
|||
|
|
print("獲取預測結果中。")
|
|||
|
|
for image_id in tqdm(image_ids):
|
|||
|
|
image_path = os.path.join(test_dir, image_id + ".png")
|
|||
|
|
img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
|
|||
|
|
|
|||
|
|
origin_shape = img.shape
|
|||
|
|
img = cv2.resize(img, (512, 512))
|
|||
|
|
|
|||
|
|
# 轉換為批次大小為 1，通道數為 1，大小為 512x512 的陣列
|
|||
|
|
img = img.reshape(1, 1, img.shape[0], img.shape[1])
|
|||
|
|
|
|||
|
|
img_tensor = torch.from_numpy(img)
|
|||
|
|
|
|||
|
|
img_tensor = img_tensor.to(device=device, dtype=torch.float32)
|
|||
|
|
|
|||
|
|
pred = net(img_tensor)
|
|||
|
|
|
|||
|
|
pred = np.array(pred.data.cpu()[0])[0]
|
|||
|
|
pred[pred >= 0.5] = 255
|
|||
|
|
pred[pred < 0.5] = 0
|
|||
|
|
pred = cv2.resize(pred, (origin_shape[1], origin_shape[0]), interpolation=cv2.INTER_NEAREST)
|
|||
|
|
cv2.imwrite(os.path.join(pred_dir, image_id + ".png"), pred)
|
|||
|
|
|
|||
|
|
print("預測結果獲取完成。")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if miou_mode == 0 or miou_mode == 2:
|
|||
|
|
print("計算 mIoU 中。")
|
|||
|
|
print(gt_dir)
|
|||
|
|
print(pred_dir)
|
|||
|
|
print(num_classes)
|
|||
|
|
print(name_classes)
|
|||
|
|
hist, IoUs, PA_Recall, Precision = compute_mIoU(gt_dir, pred_dir, image_ids, num_classes,
|
|||
|
|
name_classes) # 執行計算 mIoU 的函數
|
|||
|
|
print("mIoU 計算完成。")
|
|||
|
|
miou_out_path = "results/"
|
|||
|
|
show_results(miou_out_path, hist, IoUs, PA_Recall, Precision, name_classes)
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
cal_miou()
|