子豪兄2天完成论文系列之预测图像

1 预测单张图像

import os

import cv2

import pandas as pd
import numpy as np

import torch

import matplotlib.pyplot as plt
%matplotlib inline

设备

# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

print('device', device)

device cuda:0

载入预训练图像分类模型

import torch
from torch import nn
import torchvision
from torch.nn import functional as F
from efficientnet_pytorch import EfficientNet

class Detector(nn.Module):

    def __init__(self):
        super(Detector, self).__init__()
        self.net=EfficientNet.from_pretrained("efficientnet-b4",advprop=True,num_classes=2)
        

    def forward(self,x):
        x=self.net(x)
        return x

model=Detector()
model=model.to(device)
cnn_sd=torch.load('weights/FFraw')
model.load_state_dict(cnn_sd)
model.eval()

图像预处理

import numpy as np
import cv2
from PIL import Image
import sys
from tqdm import tqdm

def extract_face(frame,model,image_size=(380,380)):


    faces = model.predict_jsons(frame)

    if len(faces)==0:
        print('No face is detected' )
        return []

    croppedfaces=[]
    for face_idx in range(len(faces)):
        x0,y0,x1,y1=faces[face_idx]['bbox']
        bbox=np.array([[x0,y0],[x1,y1]])
        croppedfaces.append(cv2.resize(crop_face(frame,None,bbox,False,crop_by_bbox=True,only_img=True,phase='test'),dsize=image_size).transpose((2,0,1)))

    return croppedfaces

载入一张测试图片

frame = cv2.imread("")

frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

# 用 pillow 载入
from PIL import Image
img_pil = Image.open(frame)

np.array(img_pil).shape

执行分类预测

from retinaface.pre_trained_models import get_model
face_detector = get_model("resnet50_2020-07-20", max_size=max(frame.shape),device=device)
face_detector.eval()
face_list=extract_face(frame,face_detector)

face_list.shape

# img = face_list.unsqueeze(0).to(device)
# 执行前向预测，得到所有类别的 logit 预测分数
#pred_logits = model(input_img) 

# pred_logits
#import torch.nn.functional as F
#pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算

#pred_softmax.shape

with torch.no_grad():
    img=torch.tensor(face_list).to(device).float()/255
    # torchvision.utils.save_image(img, f'test.png', nrow=8, normalize=False, range=(0, 1))
    pred=model(img).softmax(1)[:,1].cpu().data.numpy().tolist()

pred.shape