使用HRnet训练自己的模型并检测

最新推荐文章于 2024-06-17 18:15:00 发布

qq_44934014

最新推荐文章于 2024-06-17 18:15:00 发布

阅读量2.2k

点赞数 2

本文链接：https://blog.csdn.net/qq_44934014/article/details/121181461

版权

HRNet 图像分类模型加载 PyTorch 预处理

关键词由CSDN通过智能技术生成

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import cv2 as cv
import torch
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
from torch.autograd import Variable
from config import config
from config import update_config
from PIL import Image
import argparse
from models import cls_hrnet



def prediect():

    # loda HRNET
    config.merge_from_file('C:\\Users\\User\\Desktop\\HRnet-class\\HRNet-Classification\\experiments\\cls_hrnet_w18_small_v2_sgd_lr5e-2_wd1e-4_bs32_x100.yaml')
    config.freeze()

    #
    # parser = argparse.ArgumentParser(description='Train network')
    # parser.add_argument('--TEST.MODEL_FILE',
    #                     help='model path',
    #                     type = str,
    #                     default='')
    # args = parser.parse_args()
    # update_config(config,args)

    cudnn.benchmark = config.CUDNN.BENCHMARK
    torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
    torch.backends.cudnn.enabled = config.CUDNN.ENABLED
    hrnet = cls_hrnet.get_cls_net(config)
    print('*********************************')
    print(config.TEST.MODEL_FILE)

    if config.TEST.MODEL_FILE:
        hrnet.load_state_dict(torch.load(config.TEST.MODEL_FILE))
    else:
        print('没找到模型文件')

    gpus = list(config.GPUS)
    hrnet = torch.nn.DataParallel(hrnet, device_ids=gpus).cuda()
    hrnet.eval()

    Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor

    pli_img_path = r'C:\Users\User\Desktop\HRnet-class\HRNet-Classification\imagenet\images\train\normal\im0088.jpg'
    pil_img = Image.open(pli_img_path)

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    input = transforms.Compose([
        transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)),
        transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]),
        transforms.ToTensor(),
        normalize,
    ])(pil_img)
    input = Variable(torch.unsqueeze(input, dim=0).float(), requires_grad=False)
    # switch to evaluate mode
    cls_pred = 0
    with torch.no_grad():
        output = hrnet(input)

        # print(output)

        # free image
        torch.cuda.empty_cache()

        cls_pred = output.argmax(dim=1)
        print(cls_pred)

        if cls_pred ==0:
            print('This pic belong to class:')
            print('fall')

        if cls_pred ==1:
            print('This pic belong to class:')
            print('normal')



prediect()