update

ieemoo-ai-searchv2.py

@@ -0,0 +1,96 @@
+import sys
+import argparse
+#from utils.retrieval_index import EvaluteMap
+from utils.tools import EvaluteMap
+from utils.retrieval_feature import AntiFraudFeatureDataset
+from utils.monitor import Moniting
+from utils.updateObs import *
+from utils.config import cfg
+from utils.tools import createNet
+from flask import request,Flask
+from utils.forsegmentation import analysis
+from gevent.pywsgi import WSGIServer
+import os, base64, stat, shutil, json, time
+sys.path.append('RAFT')
+sys.path.append('RAFT/core')
+sys.path.append('RAFT/core/utils')
+from RAFT.analysis_video import *
+import logging.config
+from skywalking import agent, config
+from threading import Thread
+
+SW_SERVER = os.environ.get('SW_AGENT_COLLECTOR_BACKEND_SERVICES')
+SW_SERVICE_NAME = os.environ.get('SW_AGENT_NAME')
+if SW_SERVER and SW_SERVICE_NAME:
+    config.init() #采集服务的地址，给自己的服务起个名称
+    #config.init(collector="123.60.56.51:11800", service='ieemoo-ai-search') #采集服务的地址，给自己的服务起个名称
+    agent.start()
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
+app = Flask(__name__)
+
+parser = argparse.ArgumentParser()
+#parser.add_argument('--model', default='../module/ieemoo-ai-search/model/now/raft-things.pth',help="restore checkpoint")
+parser.add_argument('--model', default='../module/ieemoo-ai-search/model/now/raft-small.pth',help="restore checkpoint")
+#parser.add_argument('--small', action='store_true', help='use small model')
+parser.add_argument('--small', type=bool, default=True, help='use small model')
+parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
+parser.add_argument('--alternate_corr', action='store_true', help='use efficent correlation implementation')
+opt, unknown = parser.parse_known_args()
+
+'''
+status 状态码
+00: 视频未解析成功(视频截取错误)
+01: 未纳入监查列表
+02: 未检测出商品
+03: 异常输出
+04: 正确识别
+'''
+status = ['00', '01', '02', '03', '04']
+net, transform, ms = createNet()
+raft_model = raft_init_model(opt)
+def setup_logging(path):
+    if os.path.exists(path):
+        with open(path, 'r') as f:
+            config = json.load(f)
+            logging.config.dictConfig(config)
+            logger = logging.getLogger("root")
+    return logger
+
+logger = setup_logging('utils/logging.json')
+@app.route('/searchv2', methods=['POST'])
+def search():
+    pre_status = False
+    try:
+        video_name = request.form.get('video_name')
+        logger.info('get video '+video_name)
+        ocr_file_path = os.sep.join([cfg.Ocrtxt, video_name.split('.')[0]+'.txt'])
+        video_extra_info = request.form.get('video_extra_info')
+        if not video_extra_info is None:
+           with open(ocr_file_path, 'w') as f:
+                f.write(video_extra_info)
+        video_data = request.files['video']
+        videoPath = os.sep.join([cfg.VIDEOPATH, video_name])
+        video_data.save(videoPath)
+        uuid_barcode = video_name.split('.')[0] 
+        barcode_name = uuid_barcode.split('_')[-1]
+        if Moniting(barcode_name).search() == 'nomatch':
+            state = status[1]   
+            analysis_video(raft_model, videoPath, '',uuid_barcode,None,net=net, transform=transform,ms=ms, match=False)
+        else:
+            state = analysis_video(raft_model, videoPath, '',uuid_barcode,None,net=net, transform=transform,ms=ms, match=True)
+        result = uuid_barcode+'_'+state #参数修改返回结果
+    except Exception as e:
+        logger.warning(e)   #异常返回00
+        thread = Thread(target=AddObs, kwargs={'file_path':videoPath, 'status':status[3]})
+        thread.start()
+        return uuid_barcode+'_'+status[3] #参数修改返回00   
+    thread = Thread(target=AddObs, kwargs={'file_path':videoPath, 'status':state})
+    thread.start()
+    logger.info(result)
+    print('result >>>>> {}'.format(result))
+    return result
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=8085)
+