File size: 3,619 Bytes
1847219 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
# animeGender 0.8 use
# model using file
# DOF Studio 230801
import cv2
import torch
import numpy as np
from torchvision import transforms
num_cls = 2
classes = ['female', 'male']
#############################
# graphic lib
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
#############################
# use it
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
# for those who use "image_path"
def predict_class(img_path:str, model:any, print_:bool = False):
img = cmpgraph_224x224_ret(img_path)
transform = transforms.Compose(
[
# transforms.Resize(224),
# transforms.CenterCrop(224),
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
# for those who use direct image data
def predict_img_class(img:any, model:any, print_:bool = False):
img = cmpgraph_224x224_dret(img)
transform = transforms.Compose(
[
# transforms.Resize(224),
# transforms.CenterCrop(224),
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__':
# TWO STEPS TO USE THIS MODEL
# @ DOF Studio @
# load a model from your disk
model = loadmodel("your_model_path")
# interfere an image and get the feedback
cls, confidence = predict_class("your_image_path", model, print_ = True)
|