|
- '''
- # @time:2023/3/23 20:13
- # Author:Tuan
- # @File:predict_GF.py
- '''
- '''
- # @time:2023/3/14 10:01
- # Author:Tuan
- # @File:predict_iter.py
- '''
- import os.path
- import torch.nn
- import imageio
- import numpy as np
- import cv2
- from datetime import datetime
- # from Model.HrNet.hrnet_T import HighResolutionNet
- from Model.UNet.unet import UNet
- from Model.SegFormer.segformer import SegFormer
- from Model.ResNet18.ResNet101 import resnet34
- import glob
- from utils.utils import Logger
- from datetime import datetime
- from pathlib import Path
-
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
- now = datetime.now()
- now = str(now.month) + '_' + str(now.day) + '_' + str(now.hour) + '_' + str(now.minute)
-
- def estimate(y_label, y_pred):
- # y_pred[y_label==0]=0
- # 准确率
- acc = np.mean(np.equal(y_label, y_pred))
-
- return acc, y_pred
-
- def readimage(dir):
- images_path_list = glob.glob(os.path.join(dir, '*.tiff'))
- return images_path_list
-
-
-
- def model_predict(model, img_data, img_size):
- row, col, dep = img_data.shape
-
- if row % img_size != 0 or col % img_size != 0:
- # 计算填充后图像的 hight 和 width
- padding_h = (row // img_size + 1) * img_size
- padding_w = (col // img_size + 1) * img_size
- else:
- # 不填充后图像的 hight 和 width
- padding_h = (row // img_size) * img_size
- padding_w = (col // img_size) * img_size
-
- # 初始化一个 0 矩阵,将图像的值赋值到 0 矩阵的对应位置
- padding_img = np.zeros((padding_h, padding_w, dep), dtype='float32')
- padding_img[:row, :col, :] = img_data[:row, :col, :]
-
- # 初始化一个 0 矩阵,用于将预测结果的值赋值到 0 矩阵的对应位置
- padding_pre = np.zeros((padding_h, padding_w), dtype='uint8')
- padding_pre_sl = np.zeros((padding_h, padding_w), dtype='uint8')
- padding_pre_line = np.zeros((padding_h, padding_w), dtype='uint8')
- # padding_pre = torch.tensor(padding_pre)
- # padding_pre = padding_pre.to(device)
-
- # 对 img_size * img_size 大小的图像进行预测
- count = 0 # 用于计数
- for i in list(np.arange(0, padding_h, img_size)):
- if (i + img_size) > padding_h:
- continue
- for j in list(np.arange(0, padding_w, img_size)):
- if (j + img_size) > padding_w:
- continue
-
- # 取 img_size 大小的图像,在第一维添加维度,变成四维张量,用于模型预测
- img_data_ = padding_img[i:i + img_size, j:j + img_size, :]
- img_data_ = img_data_[np.newaxis, :, :, :]
- img_data_ = np.transpose(img_data_, (0, 3, 1, 2))
- img_data_ = torch.from_numpy(img_data_).to(device)
-
- # 预测,对结果进行处理
- y_pre_seg = model.forward(img_data_)
- # y_pre = model.predict(img_data_)
- y_pre = np.squeeze(y_pre_seg, axis=0)
- y_pre = torch.argmax(y_pre, axis=0)
- # y_pre = y_pre.astype('uint8')
-
- # 将预测结果的值赋值到 0 矩阵的对应位置
- padding_pre[i:i + img_size, j:j + img_size] = y_pre[:img_size, :img_size].cpu()
-
-
- count += 1 # 每预测一块就+1
-
- # 计算准确率
- y_pre = padding_pre[:row, :col]+1
-
- return y_pre
-
- #参数
- import argparse
- parser = argparse.ArgumentParser(add_help=False)
- parser.add_argument("--modelname",type=str,default=" ")
- args = parser.parse_args()
- #参数
- # num_class = 8
- os.system("ls")
- image_size = 1024
- modelname = "GF"
- imagedir = r"/dataset"
- # labeldir = r"/dataset/coastline_classify_label_vaild"
- imagelist = readimage(imagedir)
- # labellist = readimage(labeldir)
-
- print(f"测试数量为:{len(imagelist)}")
- print(os.listdir("/dataset"))
- modelPath = r"/model/unet_GF"#r"E:\yqj\try\code\torch\Train\save_model\UNet\510-0.01232.pth"#
- if os.path.exists(modelPath):
- print(glob.glob(os.path.join(modelPath, '*.pth')))
- elif os.path.exists(r"/pretrainmodel/unet_GF"):
- modelPath = r"/pretrainmodel/unet_GF"
- print(modelPath)
- else:
- print("no")
- modelPath= glob.glob(os.path.join(modelPath, '*.pth'))
- print(modelPath)
- print(f"模型数量为:{len(modelPath)}")
-
- savePath = r"/result//"+modelname
- log_path = r"/result" + "//"+modelname+"//"+ now+ ".log"
-
- # name = Path(modelPath).stem
-
- if os.path.exists(savePath) == False:
- os.makedirs(savePath)
-
- f = open(log_path, 'w')
- f.close()
-
- log = Logger(log_path, level='debug')
- log.logger.info('Start! Train image size ' + str(image_size))
-
-
-
- # 加载模型
- # model = SegFormer(num_classes=7, phi= 'b0',pretrained=False).to(device)
- # model = resnet34(4,7)
- model = UNet(num_classes=16)
- # model.load_state_dict(torch.load(modelPath, map_location = torch.device('cpu')))
- for i in range(len(modelPath)):
- log.logger.info('modelPath: ' + modelPath[i])
- print(f'modelPath:{modelPath[i]}')
- name = Path(modelPath[i]).stem
- if os.path.exists(os.path.join (savePath,name)) == False:
- os.makedirs( os.path.join (savePath,name))
- model.load_state_dict(torch.load(modelPath[i], map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu")))
- model.eval()
- acc_all = 0
- for j in range(len(imagelist)):
- image = imageio.imread(imagelist[j])
- # label = imageio.imread(labellist[i])
- # image = Image.open(imagePath)
- # image = np.array(image)
- # label = Image.open(labelPath)
- # label = np.array(label)
- B1, B2, B3, B4 = cv2.split(image)
- B1_normalization = ((B1 - np.min(B1)) / (np.max(B1) - np.min(B1)) * 1).astype('float32')
- B2_normalization = ((B2 - np.min(B2)) / (np.max(B2) - np.min(B2)) * 1).astype('float32')
- B3_normalization = ((B3 - np.min(B3)) / (np.max(B3) - np.min(B3)) * 1).astype('float32')
- B4_normalization = ((B4 - np.min(B4)) / (np.max(B4) - np.min(B4)) * 1).astype('float32')
- image = cv2.merge([B1_normalization, B2_normalization, B3_normalization, B4_normalization])
-
- row, col = image.shape[0],image.shape[1]
- row_new = row//256
- if row%256 != 0:
- row_new += 1
- col_new = col//256
- if col%256 != 0:
- col_new += 1
- # padding_img = np.zeros((row_new*256, col_new*256, 4), dtype='float32')
- # padding_img[:row, :col, :] = image[:row, :col, :]
- # padding_pre = np.zeros((row_new*256, col_new*256), dtype='uint8')
- # padding_pre_line = np.zeros((row_new * 256, col_new * 256), dtype='uint8')
- # img_data_ = padding_img[np.newaxis, :, :, :]
- # img_data_ = np.transpose(img_data_, (0, 3, 1, 2))
- # img_data_ = torch.from_numpy(img_data_)
-
- # 预测,对结果进行处理
- # y_pre_seg,_ = model(img_data_.to(device))
- # # y_pre = model.predict(img_data_)
- # y_pre = np.squeeze(y_pre_seg, axis=0)
- # y_pre = torch.argmax(y_pre, axis=0)
- #
- # padding_pre[:row, :col] = y_pre[:row, :col].cpu()
- # acc , y_pred= estimate(label, padding_pre[:row, :col] + 1)
- y_pred = model_predict(model.to(device), image, img_size=image_size)
-
- # output = output.numpy()
- # output = output.argmax(dim = 0)
- # print(f"{i + 2016}的准确率: {acc}")
- print(y_pred.shape)
- save = os.path.join(savePath , name + "_seg.tif")
- imageio.imwrite( save,y_pred)
- # y_pre = model.forward_line(y_pre_seg.to(device),img_data_.to(device))
- # # y_pre = model.predict(img_data_)
- # y_pre = np.squeeze(y_pre[0], axis=0)
- # y_pre = torch.argmax(y_pre, axis=0)
- # padding_pre_line[:row, :col] = y_pre[:row, :col].cpu()
-
- # print("finish")
-
-
- # print(f"平均准确率: {acc_all/len(imagelist)}")
-
|