From 55962c44e4601587996136e48a1ea35e8176b95e Mon Sep 17 00:00:00 2001
From: Brainway <brainway@protonmail.com>
Date: Tue, 18 Oct 2022 03:39:44 +0000
Subject: [PATCH] update ieemoo-ai-isempty.py.

---
 ieemoo-ai-isempty.py | 35 ++++++++++++++++++++++-------------
 1 file changed, 22 insertions(+), 13 deletions(-)

diff --git a/ieemoo-ai-isempty.py b/ieemoo-ai-isempty.py
index 0cb0ae1..14a0667 100755
--- a/ieemoo-ai-isempty.py
+++ b/ieemoo-ai-isempty.py
@@ -45,12 +45,12 @@ print(torch.__version__)
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--img_size", default=600, type=int, help="Resolution size")
+    parser.add_argument("--img_size", default=320, type=int, help="Resolution size")
     parser.add_argument('--split', type=str, default='overlap', help="Split method")
     parser.add_argument('--slide_step', type=int, default=2, help="Slide step for overlap split")
     parser.add_argument('--smoothing_value', type=float, default=0.0, help="Label smoothing value")
     parser.add_argument("--pretrained_model", type=str, default="../module/ieemoo-ai-isempty/model/now/emptyjudge5_checkpoint.bin", help="load pretrained model")
-    #parser.add_argument("--pretrained_model", type=str, default="./output/ieemooempty_vit_checkpoint.pth", help="load pretrained model")
+    #parser.add_argument("--pretrained_model", type=str, default="output/ieemooempty_vit_checkpoint.pth", help="load pretrained model") #使用自定义VIT
     opt, unknown = parser.parse_known_args()
     return opt
 
@@ -65,7 +65,7 @@ class Predictor(object):
         self.num_classes = 0
         self.model = None
         self.prepare_model()
-        self.test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEAR),
+        self.test_transform = transforms.Compose([transforms.Resize((320, 320), Image.BILINEAR),
                                              transforms.ToTensor(),
                                              transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
 
@@ -102,7 +102,7 @@ class Predictor(object):
                 topN = torch.argsort(probs, dim=-1, descending=True).tolist()
                 clas_ids = topN[0][0]
                 clas_ids = 0 if 0==int(clas_ids) or 2 == int(clas_ids) or 3 == int(clas_ids) else 1
-                #print("cur_img result: class id: %d, score: %0.3f" % (clas_ids, probs[0, clas_ids].item()))
+                print("cur_img result: class id: %d, score: %0.3f" % (clas_ids, probs[0, clas_ids].item()))
                 result["success"] = "true"
                 result["rst_cls"] = str(clas_ids)
                 return result
@@ -114,21 +114,16 @@ predictor = Predictor(args)
 
 @app.route("/isempty", methods=['POST'])
 def get_isempty():
-    start = time.time()
-    #print('--------------------EmptyPredict-----------------')
+    print("begin")
+
     data = request.get_data()
-    ip = request.remote_addr
-    #print('------ ip = %s ------' % ip)
-    print(ip)
 
     json_data = json.loads(data.decode("utf-8"))
-    getdateend = time.time()
-    #print('get date use time: {0:.2f}s'.format(getdateend - start))
 
     pic = json_data.get("pic")
-    result = {}
-              
     imgdata = base64.b64decode(pic)
+    
+    result ={}
     imgdata_np = np.frombuffer(imgdata, dtype='uint8')
     img_src = cv2.imdecode(imgdata_np, cv2.IMREAD_COLOR)
     img_data = Image.fromarray(np.uint8(img_src))
@@ -136,5 +131,19 @@ def get_isempty():
    
     return repr(result)
 
+def getByte(path):
+    with open(path, 'rb') as f:
+        img_byte = base64.b64encode(f.read())
+    img_str = img_byte.decode('utf-8')
+    return img_str
+
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=8888)
+
+    # result ={}
+    # imgdata = base64.b64decode(getByte("img.jpg"))
+    # imgdata_np = np.frombuffer(imgdata, dtype='uint8')
+    # img_src = cv2.imdecode(imgdata_np, cv2.IMREAD_COLOR)
+    # img_data = Image.fromarray(np.uint8(img_src))
+    # result = predictor.normal_predict(img_data, result)
+    # print(result)