Algorithm-Repo/SourceCode/MachineLearningUtils/Classification/PytorchFramework/predict.py

import argparse
import os
import logging
import cv2
import numpy as np
import glob
import json
import time
import torch
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
from utils.confusion_matrix import draw_confusion_matrix

parser = argparse.ArgumentParser(description='PyTorch ImageNet Predict')
parser.add_argument('--model', '-m', default=os.path.join(os.getcwd(), 'checkpoints\\mobilenetv3\\CP_epoch31.pth'),
                    metavar='FILE',
                    help="Specify the file in which the model is stored")
parser.add_argument('--input', '-i', metavar='INPUT', nargs='+',
                    help='filenames of input images', default=os.path.join(os.getcwd(), "data\\test\\"))
parser.add_argument('--output', '-o', metavar='OUTPUT', nargs='+',
                    help='path of save the result', default=os.path.join(os.getcwd(), 'output\\20240515-augmentation-2\\'))
parser.add_argument('-a', '--arch', metavar='ARCH', default='efficientnet-b0',
                    help='model architecture (default: efficientnet-b0)')
parser.add_argument('--gpu', default=0, type=int,
                    help='GPU id to use.')
parser.add_argument('--image_size', default=256, type=int,
                    help='image size')

index_label = {0:'BMode', 1:'BModeBlood',2:'Pseudocolor',3:'PseudocolorBlood', 4:'Spectrogram', 5:'CEUS', 6:'SE',7:'STE',8:'FourDime'}

def predict_img(net,
                full_img,
                device,
                scale_factor=1):
    net.eval()                              # 测试时的网络特征

    img = cv2.resize(full_img, (scale_factor, scale_factor))
    if len(img.shape) == 2:
        img = np.expand_dims(img, axis=2)  # 表示在axis位置添加数据
    img = img.transpose((2, 0, 1))         # 转置:Pytorch中为[Channels, H, W]
    if img.max() > 1:
        img = img / 255
    img = torch.from_numpy(img)            # 将x转换为torch类型

    img = img.unsqueeze(0)                 # 扩充第0个维度
    img = img.to(device=device, dtype=torch.float32)
    with torch.no_grad():
        t = time.time()
        output = net(img).cpu().numpy()
        # print("时间：{}".format((time.time() - t)*1000))
        pred_label = np.argmax(output)
        score = output[0][int(pred_label)]
        
        # print("score:{}".format(output[0][int(pred_label)]))


    return pred_label, score

def find_images(root_dir):
    """
    深度遍历查找指定目录下的所有图像文件
    """
    swap_dict = {v:k for k,v in index_label.items()}
    image_files = []
    for root, dirs, files in os.walk(root_dir):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):                
                json_file = os.path.join(root, file[0:-4] + ".txt")
                labeled_info = open(json_file, 'r')
                txt_info = labeled_info.read()
                txt_info_dict = json.loads(txt_info)
                classes = txt_info_dict[0]["FileResultInfos"][0]["LabeledResult"]["ImageResults"][0]["Conclusion"]["Title"] 
                dst_label = swap_dict[classes]

                image_files.append((os.path.join(root, file), dst_label))
    return image_files

def main():
    args = parser.parse_args()

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    # EfficientNet预测
    # from model.efficientnet_pytorch.model import EfficientNet
    # net = EfficientNet.from_name(model_name=args.arch, in_channels=3, num_classes=9, image_size=args.image_size)

    # resnext50预测
    # from resnext_pytorch.resnext import resnext50
    # net = resnext50(baseWidth=4, cardinality=8)

    # mobilenetv3预测
    from model.mobilenetv3 import mobilenetv3
    net = mobilenetv3.MobileNetV3(n_class=9, input_size=256, channels=3, mode='small')

    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    device = torch.device('cpu')
    logging.info(f'Using device {device}')
    net.to(device=device)
    net.load_state_dict(torch.load(args.model, map_location=device))
    logging.info("Model loaded !")

    y_gt=[]
    y_pred=[]
    
    f_names = find_images(args.input)   # 匹配所有的符合条件的文件，并将其以list的形式返回完整路径
    times = 0
    for  i in range(len(f_names)):
        imgpath, labels_true = f_names[i]
        img = cv2.imdecode(np.fromfile(imgpath, dtype=np.uint8), 1)  # 读回数据时需要用户指定元素类型，并对数组的形状进行适当的修改
        t =time.time()
        classification_pred, score = predict_img(net=net, full_img=img, scale_factor=args.image_size, device=device)
        times += (time.time() - t)*1000
        # print("时间：{}".format((time.time() - t)*1000))

        logging.info("Visualizing results for image {}, close to continue ...".format(imgpath))

        # save results
        Label_pred_dir = os.path.join(args.output, index_label[classification_pred])
        if not os.path.exists(Label_pred_dir):
            os.makedirs(Label_pred_dir)
        # print("预测的第{}图的类别为{}".format(i, index_label[classification_pred]))
        cv2.imencode('.jpg', img.astype(np.uint8))[1].tofile(os.path.join(Label_pred_dir, imgpath.split('\\')[-1]))

        y_pred.append(classification_pred)
        y_gt.append(labels_true)

        if i%20==0:
            print(i, '/', len(f_names))
        # print("分类结果：{}\n".format(index_label[classification_pred]))

    # 绘制热力图
    draw_confusion_matrix(label_true=y_gt,
                      label_pred=y_pred,
                      label_name=[v for k, v in index_label.items()],
                      normlize=True,
                      title="Confusion Matrix",
                      pdf_save_path=os.path.join(args.output,"CM.jpg"),
                      dpi=300)


if __name__ == '__main__':
    main()