# 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)