Ver código fonte

modified: client.py
modified: server.py

meishuhuan 2 anos atrás
pai
commit
120a5faf02
5 arquivos alterados com 19 adições e 12 exclusões
  1. 3 3
      Dockerfile
  2. 2 2
      client.py
  3. 1 1
      requirements.txt
  4. 12 6
      server.py
  5. 1 0
      yolov5

+ 3 - 3
Dockerfile

@@ -93,7 +93,7 @@ startretries=0\n\
 redirect_stderr=true\n\
 stdout_logfile=/var/log/be.log\n\
 stdout_logfile_maxbytes=50MB\n\
-environment=CUDA_VISIBLE_DEVICES=0,PYTHONPATH=/worksapce/yolov5\n\
+environment=CUDA_VISIBLE_DEVICES=0,PYTHONUNBUFFERED=1,PYTHONPATH=/worksapce/yolov5\n\
 " > /etc/supervisor/conf.d/be.conf
 
 # envs/py38/
@@ -102,9 +102,9 @@ RUN git clone https://gitee.com/monkeycc/yolov5.git
 # RUN wget --user=sxkj --password='sx' ftp://192.168.199.208/yili.pt
 # RUN wget --user=sxkj --password='sx' ftp://192.168.199.208/best.pt
 # RUN wget --user=sxkj --password='sx' ftp://192.168.199.208/best2.pt
-RUN wget --user=sxkj --password='sx' ftp://192.168.199.208/yili2.pt
+RUN wget --user=sxkj --password='sx' ftp://192.168.199.208/8-14_layout_model.pt
 
-Add . /workspace
+ADD . /workspace    
 EXPOSE 8080
 
 

+ 2 - 2
client.py

@@ -13,8 +13,8 @@ from io import BytesIO
 from PIL import Image
 
 def send_request(file_list = ['./images/zidane.jpg'], 
-                    model_name = 'yolov5s',
-                    img_size = 640,
+                    model_name = 'yolov5x',
+                    img_size = 800,
                     download_image = False):
 
     #upload multiple files as list of tuples

+ 1 - 1
requirements.txt

@@ -36,4 +36,4 @@ pandas
 # extras --------------------------------------
 thop  # FLOPS computation
 pycocotools>=2.0  # COCO mAP
-torchvision>=0.8.1
+torchvision>=0.8.1

+ 12 - 6
server.py

@@ -1,3 +1,4 @@
+# -*- coding:utf-8 -*-
 from fastapi import FastAPI, Request, Form, File, UploadFile
 from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel
@@ -17,7 +18,7 @@ YOLO_DIR = '/workspace/yolov5'
 # WEIGHTS = '/workspace/best.pt'
 # WEIGHTS = '/workspace/yili.pt'
 # WEIGHTS = '/workspace/best2.pt'
-WEIGHTS = '/workspace/yili2.pt'
+WEIGHTS = '/workspace/8-14_layout_model.pt'
 
 logger = logging.getLogger('log')
 logger.setLevel(logging.DEBUG)
@@ -42,7 +43,7 @@ logger.addHandler(ch)
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 bl = torch.cuda.is_available()
-logger.info(f'是否可使用GPU=======>{bl}')
+#logger.info(f'是否可使用GPU=======>{bl}')
 
 app = FastAPI()
 templates = Jinja2Templates(directory = 'templates')
@@ -55,6 +56,7 @@ colors = [tuple([random.randint(0, 255) for _ in range(3)]) for _ in range(100)]
 
 if model_dict['ocr-layout'] is None:
           model_dict['ocr-layout'] = model = torch.hub.load(YOLO_DIR, 'custom', path=WEIGHTS, source='local').to(device)
+#          print(model_dict['ocr-layout'])
           logger.info("========>模型加载成功")
 
 
@@ -91,7 +93,7 @@ def drag_and_drop_detect(request: Request):
 def detect_via_web_form(request: Request,
                         file_list: List[UploadFile] = File(...),
                         model_name: str = Form(...),
-                        img_size: int = Form(1824)):
+                        img_size: int = Form(800)):
 
     '''
     Requires an image file upload, model name (ex. yolov5s). Optional image size parameter (Default 1824).
@@ -109,7 +111,8 @@ def detect_via_web_form(request: Request,
     #create a copy that corrects for cv2.imdecode generating BGR images instead of RGB
     #using cvtColor instead of [...,::-1] to keep array contiguous in RAM
     img_batch_rgb = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in img_batch]
-
+    print(img_size)
+    print('111111')
     results = model_dict[model_name](img_batch_rgb, size = img_size)
 
     json_results = results_to_json(results,model_dict[model_name])
@@ -139,7 +142,7 @@ def detect_via_web_form(request: Request,
 def detect_via_api(request: Request,
                 file_list: List[UploadFile] = File(...),
                 model_name: str = Form(...),
-                img_size: Optional[int] = Form(1824),
+                img_size: Optional[int] = Form(800),
                 download_image: Optional[bool] = Form(False)):
 
     '''
@@ -159,7 +162,10 @@ def detect_via_api(request: Request,
     #create a copy that corrects for cv2.imdecode generating BGR images instead of RGB,
     #using cvtColor instead of [...,::-1] to keep array contiguous in RAM
     img_batch_rgb = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in img_batch]
-
+    print('111111')
+    print(img_size)
+    img_size= 800
+    print(img_size)
     results = model_dict[model_name](img_batch_rgb, size = img_size)
     json_results = results_to_json(results,model_dict[model_name])
 

+ 1 - 0
yolov5

@@ -0,0 +1 @@
+Subproject commit 7639e4c7efc6d660fefcea5589d467e88afd8b6c