update ieemoo-ai-isempty.py.

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])])
 
@@ -82,7 +82,7 @@ class Predictor(object):
         #         self.model = torch.load(self.args.pretrained_model)
         #     else:
         #         self.model = torch.load(self.args.pretrained_model,map_location='cpu')
-        self.model = torch.load(self.args.pretrained_model)
+        self.model = torch.load(self.args.pretrained_model,map_location=torch.device('cpu'))
         self.model.eval()
         self.model.to("cuda")
         
@@ -95,14 +95,14 @@ class Predictor(object):
         else:
             with torch.no_grad():
                 x = self.test_transform(img_data)
-                if torch.cuda.is_available():
-                    x = x.cuda()
+                # if torch.cuda.is_available():
+                #     x = x.cuda()
                 part_logits = self.model(x.unsqueeze(0))
                 probs = torch.nn.Softmax(dim=-1)(part_logits)
                 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)