# This file is provided by DOF Studio on huggingface.co
# with model animeGender-dvgg-0.7.
# Link: https://huggingface.co/DOFOFFICIAL/animeGender-dvgg-0.7

import cv2
import torch
import torchvision
import numpy as np
from torchvision import transforms
from PIL import Image

num_cls = 2
classes = ['female', 'male']


def model_struct():
    model = torchvision.models.vgg16(pretrained=True) 

    last_dim = len(model.classifier) - 1
    num_fc = model.classifier[last_dim].in_features 
    
    model.classifier[last_dim] = torch.nn.Linear(num_fc, num_cls)

    model.classifier.append(torch.nn.Softmax())

    for param in model.parameters():
        param.requires_grad = False

    for param in model.classifier[last_dim].parameters():
        param.requires_grad = True
    for param in model.classifier[last_dim + 1].parameters():
        param.requires_grad = True

    model = model.cuda() #GPU

    return model
    
def cmpgraph_64x64(imgfrom, imgto):
    height, width, channels = dim(imgfrom)
    img = cv2.imread(imgfrom, 1)
    img2 = []
    if height > width:
        hnew = int(np.round(64 / width * height))
        wnew = 64
        img2 = cv2.resize(img, (wnew, hnew))
        img2 = img2[0:64, 0:64]
    elif width > height:
        wnew = int(np.round(64 / height * width))
        hnew = 64
        img2 = cv2.resize(img, (wnew, hnew))
        img2 = img2[0:64, 0:64]
    else:
        img2 = cv2.resize(img, (64,64))
    img3 = cv2.cvtColor(img2, cv2.COLOR_BGRA2BGR)
    return cv2.imwrite(imgto, img3)

def predict_class(img_path, model):
    img = Image.open(img_path)
    transform = transforms.Compose([transforms.ToTensor()])
    img = transform(img).cuda()
    img = torch.unsqueeze(img, dim=0)
    out = model(img)
    max = torch.max(out).item()
    pmax = torch.max(out, 1)[1].item()
    cls = classes[pmax]
    print('This is ' + cls + ' with confidence of ' + str(np.round(max, 3)))

def modelload(modelpath):
    device = torch.device('cuda')
    model = model_struct()
    model.to(device)
    model.eval()
    save = torch.load(modelpath)
    model.load_state_dict(save)
    return model

def predictmain(model, filepath):
    img = filepath
    predict_class(img, model)


if __name__ == '__main__':
    # local usage
        model = modelload("pytorch_model.bin")

    # or uses safetensors
    # model = torch.load("pytorch_model.safetensors")
    # model = model.to(torch.device('cuda'))
    
    # use your picture to interfere
    cmpgraph_64x64("path.png", "path(1).png")
    predictmain(model, "path(1).png")