Upload use.py (#2)

- Upload use.py (0833d3b92630341a7ef63e524f25f2d3632cf4de)


Co-authored-by: NathMath huang <[email protected]>

use/use.py

@@ -0,0 +1,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)
+