lichen/ieemoo-ai-isempty
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update ieemoo-ai-isempty.py.

Browse Source
This commit is contained in:
Brainway
2022-10-26 02:05:26 +00:00
committed by Gitee
parent e558921a6b
commit a699574975
1 changed files with 16 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
ieemoo-ai-isempty.py
View File

@ -45,12 +45,12 @@ print(torch.__version__)
def parse_args(): def parse_args():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--img_size", default=320, type=int, help="Resolution size") parser.add_argument("--img_size", default=600, type=int, help="Resolution size")
parser.add_argument('--split', type=str, default='overlap', help="Split method") 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('--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('--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="../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") #使用自定义VIT #parser.add_argument("--pretrained_model", type=str, default="./output/ieemooempty_vit_checkpoint.pth", help="load pretrained model")
opt, unknown = parser.parse_known_args() opt, unknown = parser.parse_known_args()
return opt return opt
@ -65,7 +65,7 @@ class Predictor(object):
self.num_classes = 0 self.num_classes = 0
self.model = None self.model = None
self.prepare_model() self.prepare_model()
self.test_transform = transforms.Compose([transforms.Resize((320, 320), Image.BILINEAR), self.test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEAR),
transforms.ToTensor(), transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) 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) # self.model = torch.load(self.args.pretrained_model)
# else: # else:
# self.model = torch.load(self.args.pretrained_model,map_location='cpu') # self.model = torch.load(self.args.pretrained_model,map_location='cpu')
self.model = torch.load(self.args.pretrained_model,map_location=torch.device('cpu')) self.model = torch.load(self.args.pretrained_model)
self.model.eval() self.model.eval()
self.model.to("cuda") self.model.to("cuda")
@ -95,14 +95,14 @@ class Predictor(object):
else: else:
with torch.no_grad(): with torch.no_grad():
x = self.test_transform(img_data) x = self.test_transform(img_data)
# if torch.cuda.is_available(): if torch.cuda.is_available():
# x = x.cuda() x = x.cuda()
part_logits = self.model(x.unsqueeze(0)) part_logits = self.model(x.unsqueeze(0))
probs = torch.nn.Softmax(dim=-1)(part_logits) probs = torch.nn.Softmax(dim=-1)(part_logits)
topN = torch.argsort(probs, dim=-1, descending=True).tolist() topN = torch.argsort(probs, dim=-1, descending=True).tolist()
clas_ids = topN[0][0] 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 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["success"] = "true"
result["rst_cls"] = str(clas_ids) result["rst_cls"] = str(clas_ids)
return result return result
@ -114,16 +114,21 @@ predictor = Predictor(args)
@app.route("/isempty", methods=['POST']) @app.route("/isempty", methods=['POST'])
def get_isempty(): def get_isempty():
print("begin") start = time.time()
#print('--------------------EmptyPredict-----------------')
data = request.get_data() data = request.get_data()
ip = request.remote_addr
#print('------ ip = %s ------' % ip)
print(ip)
json_data = json.loads(data.decode("utf-8")) 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") pic = json_data.get("pic")
imgdata = base64.b64decode(pic)
result = {} result = {}
imgdata = base64.b64decode(pic)
imgdata_np = np.frombuffer(imgdata, dtype='uint8') imgdata_np = np.frombuffer(imgdata, dtype='uint8')
img_src = cv2.imdecode(imgdata_np, cv2.IMREAD_COLOR) img_src = cv2.imdecode(imgdata_np, cv2.IMREAD_COLOR)
img_data = Image.fromarray(np.uint8(img_src)) img_data = Image.fromarray(np.uint8(img_src))
@ -131,19 +136,5 @@ def get_isempty():
return repr(result) 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__": if __name__ == "__main__":
app.run(host='0.0.0.0', port=8888) 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)