import sys sys.path.append('/root/data') import argparse import time from pathlib import Path from PIL import Image import numpy import numpy as np import os,base64 import cv2 import torch import logging import torch.backends.cudnn as cudnn from gevent.pywsgi import WSGIServer from numpy import random import json from models.experimental import attempt_load from utils.datasets import LoadStreams, LoadImages from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \ scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path from utils.plots import plot_one_box from utils.torch_utils import select_device, load_classifier, time_synchronized from flask import request, Flask, jsonify, request #from ocr_system.predict_all import * import logging app = Flask(__name__) # 日志系统配置 handler = logging.FileHandler('app.log', encoding='UTF-8') handler.setLevel(logging.DEBUG) logging_format = logging.Formatter( '%(asctime)s - %(levelname)s - %(filename)s - %(funcName)s - %(lineno)s - %(message)s') handler.setFormatter(logging_format) 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('--Form', type=str, default='') parser.add_argument('--source', type=str, default='data/images', 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') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--view-img', action='store_true', help='display results') parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') parser.add_argument('--nosave', action='store_true', help='do not save images/videos') parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--update', action='store_true', help='update all models') parser.add_argument('--project', default='runs/detect', help='save results to project/name') parser.add_argument('--name', default='exp', help='save results to project/name') parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') #opt = parser.parse_args() opt, unknown = parser.parse_known_args() def loadModel(): all_model = {} 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) return all_model def CheckRange(code): with open('goodsCategory.json', 'r', encoding='utf-8') as jsonPath: loadDict = json.load(jsonPath) for key in loadDict: if code in loadDict[key]: return key, loadDict[key] #return key, code return 'nonentity',None def detect(names, save_img=False, model=None): source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size save_img = not opt.nosave and not source.endswith('.txt') # save inference images webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith( ('rtsp://', 'rtmp://', 'http://')) # Initialize set_logging() device = select_device(opt.device) half = device.type != 'cpu' # half precision only supported on CUDA # Load model #model = attempt_load(weights, map_location=device) # load FP32 model stride = int(model.stride.max()) # model stride imgsz = check_img_size(imgsz, s=stride) # check img_size if half: model.half() # to FP16 # Second-stage classifier classify = False if classify: modelc = load_classifier(name='resnet101', n=2) # initialize modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval() # Set Dataloader vid_path, vid_writer = None, None if webcam: view_img = check_imshow() cudnn.benchmark = True # set True to speed up constant image size inference dataset = LoadStreams(source, img_size=imgsz, stride=stride) else: dataset = LoadImages(source, img_size=imgsz, stride=stride) # Get names and colors #names = model.module.names if hasattr(model, 'module') else model.names names = names colors = [[random.randint(0, 255) for _ in range(3)] for _ in names] # Run inference if device.type != 'cpu': model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once t0 = time.time() for path, img, im0s, vid_cap in dataset: img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() # uint8 to fp16/32 img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference t1 = time_synchronized() pred = model(img, augment=opt.augment)[0] # Apply NMS pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms) t2 = time_synchronized() # Apply Classifier if classify: pred = apply_classifier(pred, modelc, img, im0s) # Process detections for i, det in enumerate(pred): # detections per image if webcam: # batch_size >= 1 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count else: p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0) p = Path(p) # to Path s += '%gx%g ' % img.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh result = {} if len(det) == 0: result["TargetDetect"] = 'Search fail' result["tag"] = False return result if len(det): # Rescale boxes from img_size to im0 size det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() # Print results for c in det[:, -1].unique(): n = (det[:, -1] == c).sum() # detections per class s+= "{n} {names[int(c)]}{'s' * (n > 1)}, "# add to string # return results result = {} res = [] for nu , data in enumerate(reversed(det.cpu()).numpy()): res.append( {"xy1":(str(data[0]),str(data[1])), "xy2":(str(data[2]),str(data[3])), "precision":str(data[4]), "Class":names[int(data[5])]}) result["TargetDetect"] = res result["tag"] = True print(result) return result all_model = loadModel() #barcode, ocr = loadOcrModel() @app.route('/predict',methods=['POST']) def predict(): if request.method == 'POST': try: data = request.get_data() data_str = data.decode("utf-8") json_data = json.loads(data_str) pic = json_data.get("file") code = json_data.get("code") imgdata = base64.b64decode(pic) imgdata_np = np.frombuffer(imgdata, dtype='uint8') image = cv2.imdecode(imgdata_np, cv2.IMREAD_COLOR) section,names = CheckRange(code) #print('section',section,names) if section == 'nonentity' or section not in all_model.keys(): model = None section = 'nonentity' else: model = all_model[section] cv2.imwrite(os.sep.join(['data/images',code+'.jpg']), image) with torch.no_grad(): if section == 'nonentity': tar_res = {'tag':False} tar_res['TargetDetect'] = 'No goods' else: tar_res = detect(names, model=model) #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'])) 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()