|
|
|
|
|
|
|
|
|
import cv2 |
|
import torch |
|
import numpy as np |
|
from torchvision import transforms |
|
|
|
num_cls = 2 |
|
classes = ['female', 'male'] |
|
|
|
|
|
|
|
def cmpgraph_224x224_ret(imgpath:str): |
|
img = cv2.imread(imgpath, 1) |
|
height, width, channels = img.shape |
|
img2 = [] |
|
if height > width: |
|
hnew = int(np.round(224 / width * height)) |
|
wnew = 224 |
|
img2 = cv2.resize(img, (wnew, hnew), interpolation = cv2.INTER_LANCZOS4) |
|
img2 = img2[0:224, 0:224] |
|
elif width > height: |
|
wnew = int(np.round(224 / height * width)) |
|
hnew = 224 |
|
img2 = cv2.resize(img, (wnew, hnew), interpolation = cv2.INTER_LANCZOS4) |
|
img2 = img2[0:224, 0:224] |
|
elif width == 224 and height == 224: |
|
img2 = img |
|
else: |
|
img2 = cv2.resize(img, (224, 224), interpolation = cv2.INTER_LANCZOS4) |
|
img3 = cv2.cvtColor(img2, cv2.COLOR_BGRA2BGR) |
|
return img3 |
|
|
|
def cmpgraph_224x224_dret(img:any): |
|
height, width, channels = img.shape |
|
img2 = [] |
|
if height > width: |
|
hnew = int(np.round(224 / width * height)) |
|
wnew = 224 |
|
img2 = cv2.resize(img, (wnew, hnew), interpolation = cv2.INTER_LANCZOS4) |
|
img2 = img2[0:224, 0:224] |
|
elif width > height: |
|
wnew = int(np.round(224 / height * width)) |
|
hnew = 224 |
|
img2 = cv2.resize(img, (wnew, hnew), interpolation = cv2.INTER_LANCZOS4) |
|
img2 = img2[0:224, 0:224] |
|
elif width == 224 and height == 224: |
|
img2 = img |
|
else: |
|
img2 = cv2.resize(img, (224, 224), interpolation = cv2.INTER_LANCZOS4) |
|
img3 = cv2.cvtColor(img2, cv2.COLOR_BGRA2BGR) |
|
return img3 |
|
|
|
|
|
|
|
def loadmodel(model_path:str, is_cuda:bool=True): |
|
model = torch.load(model_path) |
|
if is_cuda == True: |
|
model.to(torch.device('cuda')) |
|
else: |
|
model.to(torch.device('cpu')) |
|
return model |
|
|
|
|
|
def predict_class(img_path:str, model:any, print_:bool = False): |
|
img = cmpgraph_224x224_ret(img_path) |
|
transform = transforms.Compose( |
|
[ |
|
|
|
|
|
transforms.ToTensor() |
|
]) |
|
img = transform(img).cuda() |
|
img = torch.unsqueeze(img, dim=0) |
|
model.eval() |
|
out = model(img) |
|
out = torch.nn.functional.softmax(out) |
|
max = torch.max(out).item() |
|
pmax = torch.max(out, 1)[1].item() |
|
cls = classes[pmax] |
|
if print_ == True: |
|
print('This is ' + cls + ' with a confidence of ' + str(np.round(max, 3))) |
|
return cls, max |
|
|
|
|
|
def predict_img_class(img:any, model:any, print_:bool = False): |
|
img = cmpgraph_224x224_dret(img) |
|
transform = transforms.Compose( |
|
[ |
|
|
|
|
|
transforms.ToTensor() |
|
]) |
|
img = transform(img).cuda() |
|
img = torch.unsqueeze(img, dim=0) |
|
model.eval() |
|
out = model(img) |
|
out = torch.nn.functional.softmax(out) |
|
max = torch.max(out).item() |
|
pmax = torch.max(out, 1)[1].item() |
|
cls = classes[pmax] |
|
if print_ == True: |
|
print('This is ' + cls + ' with a confidence of ' + str(np.round(max, 3))) |
|
return cls, max |
|
|
|
if __name__ == '__main__': |
|
|
|
|
|
|
|
|
|
|
|
model = loadmodel("your_model_path") |
|
|
|
|
|
cls, confidence = predict_class("your_image_path", model, print_ = True) |
|
|