修改部署区分cpu gpu

Dockerfile

@@ -55,10 +55,31 @@ RUN ln /opt/conda/bin/mamba /usr/local/bin/mamba && mamba init zsh
 
 FROM builder1 as builder2
 
+RUN apt-get update && apt-get install -y --no-install-recommends openssh-server && rm -rf /var/lib/apt/lists/*
+RUN mkdir /var/run/sshd
+RUN echo 'root:root' | chpasswd
+RUN sed -i 's/.*PermitRootLogin .*/PermitRootLogin yes/' /etc/ssh/sshd_config
+# SSH login fix. Otherwise user is kicked off after login
+RUN sed -i 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' /etc/pam.d/sshd
+
+RUN echo "\
+[program:sshd] \n\
+command=/usr/sbin/sshd -D\n\
+autorestart=True\n\
+autostart=True\n\
+redirect_stderr = true\n\
+" > /etc/supervisor/conf.d/sshd.conf
+
+EXPOSE 22
+
+
+FROM builder2 as builder3
+
 WORKDIR /workspace
 ADD environment.yml /environment.yml
 RUN mamba update -n base -c defaults conda -y && mamba env create -f /environment.yml && rm -rf /root/.cache
 
+
 # RUN /opt/conda/envs/py38/bin/python -m ipykernel install --name py38 --display-name "py38"
 # RUN echo "c.MultiKernelManager.default_kernel_name = 'py38'">>/root/.jupyter/jupyter_notebook_config.py
 RUN echo "\
@@ -72,29 +93,12 @@ stdout_logfile=/var/log/be.log\n\
 stdout_logfile_maxbytes=0\n\
 " > /etc/supervisor/conf.d/be.conf
 
+ARG VERSION
+ENV USE_CUDA $VERSION
+
 Add . /workspace
 EXPOSE 8080
 
-FROM builder2 as builder3
-
-RUN apt-get update && apt-get install -y --no-install-recommends openssh-server && rm -rf /var/lib/apt/lists/*
-RUN mkdir /var/run/sshd
-RUN echo 'root:root' | chpasswd
-RUN sed -i 's/.*PermitRootLogin .*/PermitRootLogin yes/' /etc/ssh/sshd_config
-# SSH login fix. Otherwise user is kicked off after login
-RUN sed -i 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' /etc/pam.d/sshd
-
-RUN echo "\
-[program:sshd] \n\
-command=/usr/sbin/sshd -D\n\
-autorestart=True\n\
-autostart=True\n\
-redirect_stderr = true\n\
-" > /etc/supervisor/conf.d/sshd.conf
-
-EXPOSE 22
-
-
 
 # RUN mamba install -y jupyterlab -n base && mamba init zsh
 # RUN /opt/conda/bin/jupyter notebook --generate-config && \

Makefile

@@ -7,5 +7,5 @@ AUTHOR          := $(shell git show -s --format='%an')
 
 .PHONY: local
 local:
-	@docker build -t registry.cn-hangzhou.aliyuncs.com/sxtest/$(NAME):$(VERSION) .
+	@docker build -t registry.cn-hangzhou.aliyuncs.com/sxtest/$(NAME):$(VERSION) --build-arg VERSION=$(VERSION) .
 	@docker push registry.cn-hangzhou.aliyuncs.com/sxtest/$(NAME):$(VERSION)

core/ocr.py

@@ -56,7 +56,7 @@ class IdCardOcr:
         return {
             "confidence": conf,
             "card_type": image_type,
-            "orientation": angle // 90,
+            "orientation": (4 - angle // 90) % 4,  # 原angle是逆时针，转成顺时针
             "name": content["Name"].to_dict(),
             "id": content["IDNumber"].to_dict(),
             "ethnicity": content["Nationality"].to_dict(),

docker-compose.yml

@@ -4,7 +4,7 @@ services:
     hostname: idcard
     container_name: idcard
     restart: always
-    image: idcard:gpu
+    image: registry.cn-hangzhou.aliyuncs.com/sxtest/idcard:cpu
     privileged: true
     ipc: host
     tty: true
@@ -14,8 +14,8 @@ services:
       - '18222:22'
     volumes:
       - ./:/workspace
-    deploy:
-      resources:
-        reservations:
-          devices:
-            - capabilities: [gpu]
+#    deploy:
+#      resources:
+#        reservations:
+#          devices:
+#            - capabilities: [gpu]

server.py

@@ -6,6 +6,7 @@ from sx_utils.sximage import *
 from sx_utils.sxtime import sxtimeit
 from sx_utils.sxweb import web_try
 from core.ocr import IdCardOcr
+import os
 
 
 app = FastAPI()
@@ -21,6 +22,11 @@ app.add_middleware(
 )
 # templates = Jinja2Templates(directory='templates')
 
+use_gpu = False
+if os.getenv('USE_CUDA') == 'gpu':
+    use_gpu = True
+
+print(f'use gpu: {use_gpu}')
 
 # 初始化ocr模型和后处理模型
 ocr = PaddleOCR(use_angle_cls=True,
@@ -30,7 +36,7 @@ ocr = PaddleOCR(use_angle_cls=True,
                 rec_algorithm='CRNN',
                 ocr_version='PP-OCRv2',
                 rec_char_dict_path="./ppocr_keys_v1.txt", lang="ch",
-                use_gpu=True,
+                use_gpu=use_gpu,
                 warmup=True)
 
 

testing/ocr_test.py

@@ -22,8 +22,9 @@ class TestIdCardOcr(unittest.TestCase):
         self.assertEqual(r['status'], '000', 'status case error')
         self.assertEqual(r['result']['orientation'], 2, 'orientation case error')
 
-    def test_front_90(self):
-        image_path = './images/front-90.png'
+    def test_front_270(self):
+        image_path = './images/front-270.png'
         r = send_request(image_path)
         self.assertEqual(r['status'], '000', 'status case error')
-        self.assertEqual(r['result']['orientation'], 1, 'orientation case error')
+        self.assertEqual(r['result']['orientation'], 3, 'orientation case error')
+