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

update

Browse Source
This commit is contained in:
lichen
2022-04-12 15:35:17 +08:00
parent c30cdd9557
commit 1ea6546e85
2 changed files with 28 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
ieemoo-ai-predict.py
View File

@ -37,9 +37,9 @@ app.logger.addHandler(handler)
load model
'''
parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default='weightsv5s', help='model.pt path(s)')
parser.add_argument('--weights', nargs='+', type=str, default='../module/ieemoo-ai-predict/', help='model.pt path(s)')
parser.add_argument('--Form', type=str, default='')
parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
parser.add_argument('--source', type=str, default='../module/ieemoo-ai-predict/imgs', help='source') # file/folder, 0 for webcam
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
@ -65,7 +65,8 @@ def loadModel():
device = select_device(opt.device)
half = device.type != 'cpu'
for name in os.listdir(opt.weights):
all_model[name.split('.')[0]] = attempt_load(os.sep.join(['weightsv5s', name]), map_location=device)
if '.pt' in name:
all_model[name.split('.')[0]] = attempt_load(os.sep.join([opt.weights, name]), map_location=device)
return all_model
def CheckRange(code):
@ -177,6 +178,7 @@ def predict():
if request.method == 'POST':
try:
data = request.get_data()
print(type(data))
data_str = data.decode("utf-8")
json_data = json.loads(data_str)
pic = json_data.get("file")
@ -191,7 +193,7 @@ def predict():
section = 'nonentity'
else:
model = all_model[section]
cv2.imwrite(os.sep.join(['data/images',code+'.jpg']), image)
cv2.imwrite(os.sep.join(['../module/ieemoo-ai-predict/imgs',code+'.jpg']), image)
with torch.no_grad():
if section == 'nonentity':
tar_res = {'tag':False}
@ -201,11 +203,11 @@ def predict():
#ocr_res = detectOcr(barcode, ocr, 'data/images/'+code+'.jpg')
ocr_res = {"Ocr": {"barcode_rst": "", "barcode_success": "false", "ocr_rst": "", "ocr_success": False}}
ocr_res.update(tar_res)
os.remove(os.sep.join(['data/images',code+'.jpg']))
os.remove(os.sep.join(['../module/ieemoo-ai-predict/imgs',code+'.jpg']))
return ocr_res
except:
return {"Ocr": {"barcode_rst": "", "barcode_success": "false", "ocr_rst": "", "ocr_success": False}, "TargetDetect": {""}, "tag": False}
if __name__ == '__main__':
#http_server = WSGIServer(('0.0.0.0', 8084), app)
#http_server.serve_forever()
app.run()
http_server = WSGIServer(('0.0.0.0', 8000), app)
http_server.serve_forever()
#app.run()