Support for ARM devices (#55)

* Support for arm devices with updated OCR

* Updated Models Download Link

api/config/config.go

@@ -65,11 +65,11 @@ type (
 		ClassificationProvider string `envconfig:"SMRITI_ML_CLASSIFICATION_PROVIDER" default:"pytorch"`
 		ClassificationParams   string `envconfig:"SMRITI_ML_CLASSIFICATION_PARAMS" default:"{\"file\":\"classification_v20230731.pt\"}"`
 		OCRProvider            string `envconfig:"SMRITI_ML_OCR_PROVIDER" default:"paddlepaddle"`
-		OCRParams              string `envconfig:"SMRITI_ML_OCR_PARAMS" default:"{\"det_model_dir\":\"det_infer\",\"rec_model_dir\":\"rec_infer\",\"cls_model_dir\":\"cls_infer\"}"`
+		OCRParams              string `envconfig:"SMRITI_ML_OCR_PARAMS" default:"{\"det_model_dir\":\"det_onnx\",\"rec_model_dir\":\"rec_onnx\",\"cls_model_dir\":\"cls_onnx\"}"`
 		SearchProvider         string `envconfig:"SMRITI_ML_SEARCH_PROVIDER" default:"pytorch"`
 		SearchParams           string `envconfig:"SMRITI_ML_SEARCH_PARAMS" default:"{\"tokenizer_dir\":\"search_tokenizer\",\"processor_dir\":\"search_processor\",\"text_file\":\"search_text_v20230731.pt\",\"vision_file\":\"search_vision_v20230731.pt\"}"` //nolint:lll
 		FacesProvider          string `envconfig:"SMRITI_ML_FACES_PROVIDER" default:"pytorch"`
-		FacesParams            string `envconfig:"SMRITI_ML_FACES_PARAMS" default:"{\"minutes\":\"1\",\"face_threshold\":\"0.9\",\"model\":\"vggface2\",\"clustering\":\"ngt\"}"`
+		FacesParams            string `envconfig:"SMRITI_ML_FACES_PARAMS" default:"{\"minutes\":\"1\",\"face_threshold\":\"0.9\",\"model\":\"vggface2\",\"clustering\":\"annoy\"}"`
 		MetadataParams         string `envconfig:"SMRITI_ML_METADATA_PARAMS" default:"{\"thumbnail_size\":\"512\"}"`
 	}
 

ml/.gitignore

@@ -2,6 +2,7 @@
 *.pdmodel
 *.pdiparams
 *.pdiparams.info
+*.txt
 __pycache__/
 *.py[cod]
 *$py.class
@@ -51,4 +52,5 @@ venv.bak/
 Thumbs.db
 .DS_Store
 .idea/
-*checkpoints*
+*checkpoints*
+*.onnx

ml/ocr/paddlepaddle.py

@@ -3,8 +3,10 @@
 import os
 import urllib.request
 import tarfile
+import subprocess
+import shutil
 
-from paddleocr import PaddleOCR
+import cv2
 
 
 MODEL_BASE_URL='https://paddleocr.bj.bcebos.com/'
@@ -28,16 +30,56 @@ def download_and_save():
     os.rename('ch_ppocr_mobile_v2.0_cls_infer', 'cls_infer')
     os.remove('cls_infer.tar')
 
+    commands = [
+        "paddle2onnx --model_dir det_infer "
+        "--model_filename inference.pdmodel "
+        "--params_filename inference.pdiparams "
+        "--save_file det_onnx/model.onnx "
+        "--opset_version 10 "
+        "--input_shape_dict=\"{'x':[-1,3,-1,-1]}\" "
+        "--enable_onnx_checker True",
+
+        "paddle2onnx --model_dir rec_infer "
+        "--model_filename inference.pdmodel "
+        "--params_filename inference.pdiparams "
+        "--save_file rec_onnx/model.onnx "
+        "--opset_version 10 "
+        "--input_shape_dict=\"{'x':[-1,3,-1,-1]}\" "
+        "--enable_onnx_checker True",
+
+        "paddle2onnx --model_dir cls_infer "
+        "--model_filename inference.pdmodel "
+        "--params_filename inference.pdiparams "
+        "--save_file cls_onnx/model.onnx "
+        "--opset_version 10 "
+        "--input_shape_dict=\"{'x':[-1,3,-1,-1]}\" "
+        "--enable_onnx_checker True"
+    ]
+    for command in commands:
+        subprocess.run(command, shell=True)
+    shutil.rmtree('det_infer')
+    shutil.rmtree('rec_infer')
+    shutil.rmtree('cls_infer')
+    urllib.request.urlretrieve('https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.7/ppocr/utils/en_dict.txt', 'rec_onnx/en_dict.txt')
+
 def load_and_run(sample):
     """Loads the saved paddle ocr models and runs sample image"""
     print('loading and running paddle ocr models')
-    ocr = PaddleOCR(use_angle_cls=True, lang='en', det_model_dir='./det_infer', rec_model_dir='./rec_infer', cls_model_dir='./cls_infer')
-    result = ocr.ocr(sample)
-    words = []
-    for res in result:
-        for line in res:
-            words.append(line[1])
-    print(words)
+    import fastdeploy as fd
+    default_option = fd.RuntimeOption()
+    default_option.use_ort_backend()
+    det_model = fd.vision.ocr.DBDetector(
+            model_file='det_onnx/model.onnx', 
+            runtime_option=default_option, model_format=fd.ModelFormat.ONNX)
+    cls_model = fd.vision.ocr.Classifier(
+        model_file='cls_onnx/model.onnx', 
+        runtime_option=default_option, model_format=fd.ModelFormat.ONNX)
+    rec_model = fd.vision.ocr.Recognizer(
+        model_file='rec_onnx/model.onnx', 
+        label_path='rec_infer/en_dict.txt', runtime_option=default_option, model_format=fd.ModelFormat.ONNX)
+    ocr = fd.vision.ocr.PPOCRv3(det_model=det_model, cls_model=cls_model, rec_model=rec_model)
+    result = ocr.predict(cv2.imread(sample))
+    print(result.text)
 
 if __name__ == '__main__':
     args = sys.argv

scripts/setup_models.py

@@ -3,7 +3,7 @@
 import requests
 import zipfile
 
-DOWNLOAD_MODELS_URL='https://www.dropbox.com/scl/fi/czw3bz5w8mcioqu393xkm/models.zip?rlkey=l1olp762n7iu3qmh4ie52ic3w&dl=1'
+DOWNLOAD_MODELS_URL='https://www.dropbox.com/scl/fi/3ee7svllbfif1uvj2p5au/models.zip?rlkey=rbdj82ptlkshw26wsc7kh9vd7&dl=1'
 
 
 print('ℹ️ downloading models, hang on...')

worker/.pylintrc

@@ -2,6 +2,6 @@
 ignore=tests
 ignore-patterns=.*_pb2.py,.*_grpc.py,.*.txt,.*file,.*.toml,.*.xml,.*.md
 disable=too-few-public-methods,logging-fstring-interpolation,broad-except,import-outside-toplevel
-extension-pkg-allow-list=ngtpy
+generated-members=fastdeploy.*,cv2.*
 [FORMAT]
 max-line-length=120

worker/Dockerfile

@@ -1,8 +1,11 @@
-FROM python:3.10-slim
+FROM python:3.10-slim AS builder
 RUN apt update && \
     apt install --no-install-recommends -y build-essential \
-    curl gcc libssl-dev python3-opencv ffmpeg libmagickwand-dev libimage-exiftool-perl exiftool libraw-dev && \
-    apt clean && rm -rf /var/lib/apt/lists/*
+    curl gcc libssl-dev python3-opencv ffmpeg libmagickwand-dev \
+    libimage-exiftool-perl exiftool libraw-dev
+RUN curl --proto '=https' --tlsv1.3 -sSf https://sh.rustup.rs -o rustup-init.sh
+RUN sh rustup-init.sh -y
+ENV PATH="/root/.cargo/bin:${PATH}"
 ADD requirements.txt /requirements.txt
 RUN pip3 install --no-cache-dir -r requirements.txt
 COPY src /app/src

worker/requirements.txt

@@ -1,20 +1,22 @@
 annoy
 asyncio
+fastdeploy
 grpcio
 grpcio-tools
 facenet-pytorch
+fastdeploy-python -f https://www.paddlepaddle.org.cn/whl/fastdeploy.html
 minio
 moviepy
-ngt
-paddlepaddle==2.4.2 -i https://pypi.tuna.tsinghua.edu.cn/simple
-paddleocr
+onnxruntime
+opencv-python
+paddle2onnx
 Pillow
 prometheus-client
 PyExifTool
 rawpy
 requests
 schedule
-Wand
 torch
 torchvision
-transformers
+transformers
+Wand

worker/src/components/faces.py

@@ -5,7 +5,6 @@
 import schedule
 from google.protobuf.empty_pb2 import Empty   # pylint: disable=no-name-in-module
 from annoy import AnnoyIndex
-import ngtpy
 
 from src.protos.api_pb2_grpc import APIStub
 from src.protos.api_pb2 import (  # pylint: disable=no-name-in-module
@@ -56,8 +55,6 @@ def cluster(self) -> None:
                 # build tree/index
                 if self.clustering_framework == 'annoy':
                     self._build_annoy(mediaitem_face_embeddings)
-                elif self.clustering_framework == 'ngt':
-                    self._build_ngt(mediaitem_face_embeddings)
                 # select items which need to be clustered
                 mediaitems_to_cluster = {}
                 new_cluster_idx = -1
@@ -130,22 +127,9 @@ def _build_annoy(self, mediaitem_face_embeddings: list[MediaItemFaceEmbedding]):
             self.data.add_item(i, mediaitem_face_embedding.embedding.embedding)
         self.data.build(len(mediaitem_face_embeddings))
 
-    def _build_ngt(self, mediaitem_face_embeddings: list[MediaItemFaceEmbedding]):
-        """Cluster faces using NGT library"""
-        # build index
-        ngtpy.create(b'index', len(mediaitem_face_embeddings[0].embedding.embedding))
-        self.data = ngtpy.Index(b'index')
-        for _, mediaitem_face_embedding in enumerate(mediaitem_face_embeddings):
-            self.data.insert(mediaitem_face_embedding.embedding.embedding)
-        self.data.build_index()
-        self.data.save()
-
     def _get_nn_for_mediaitem(self, mediaitem_idx: int):
         """Get nearest neighbour index for mediaitem"""
         if self.clustering_framework == 'annoy':
             nn_idx, nn_dist = self.data.get_nns_by_item(mediaitem_idx, n=2, include_distances=True)
             return (nn_idx[1], nn_dist[1])
-        if self.clustering_framework == 'ngt':
-            results = self.data.search(self.data.get_object(mediaitem_idx), size=2)
-            return results[1]
         return None

worker/src/providers/ocr/paddlepaddle.py

@@ -1,38 +1,55 @@
 """OCR: Paddle"""
 import logging
 
+import cv2
 from moviepy.editor import VideoFileClip
-from paddleocr import PaddleOCR
+import fastdeploy as fd
 
 
-logging.getLogger().handlers = []
-
 class PaddleModule:
     """PaddleModule OCR"""
 
     def __init__(self, params: dict) -> None:
-        self.model = PaddleOCR(show_log=False, use_angle_cls=True, lang='en',
-                               det_model_dir=f'/models/ocr/{params["det_model_dir"]}',
-                               rec_model_dir=f'/models/ocr/{params["rec_model_dir"]}',
-                               cls_model_dir=f'/models/ocr/{params["cls_model_dir"]}')
+        self.default_option = fd.RuntimeOption()
+        self.default_option.use_ort_backend()
+        self.det_model = fd.vision.ocr.DBDetector(
+            model_file=f'/models/ocr/{params["det_model_dir"]}/model.onnx',
+            runtime_option=self.default_option,
+            model_format=fd.ModelFormat.ONNX,
+        )
+        self.cls_model = fd.vision.ocr.Classifier(
+            model_file=f'/models/ocr/{params["cls_model_dir"]}/model.onnx',
+            runtime_option=self.default_option,
+            model_format=fd.ModelFormat.ONNX,
+        )
+        self.rec_model = fd.vision.ocr.Recognizer(
+            model_file=f'/models/ocr/{params["rec_model_dir"]}/model.onnx',
+            label_path=f'/models/ocr/{params["rec_model_dir"]}/en_dict.txt',
+            runtime_option=self.default_option,
+            model_format=fd.ModelFormat.ONNX,
+        )
+        self.model = fd.vision.ocr.PPOCRv3(
+            det_model=self.det_model,
+            cls_model=self.cls_model,
+            rec_model=self.rec_model,
+        )
 
     def extract(self, mediaitem_user_id: str, mediaitem_id: str, mediaitem_type: str, input_file: str) -> dict:
         """Extract text from mediaitem"""
         result, words = [], []
         if mediaitem_type == 'photo':
-            result = self.model.ocr(input_file)
+            result = self.model.predict(cv2.imread(input_file))
+            if result is not None:
+                words = result.text
         else:
             video_clip = VideoFileClip(input_file)
             for frame in video_clip.iter_frames(fps=video_clip.fps):
-                _result = self.model.ocr(frame)
-                result += _result
+                _result = self.model.predict(frame)
+                if _result is not None:
+                    result += _result.text
             video_clip.reader.close()
-        for res in result:
-            for line in res:
-                if line[1][1] > 0.90:
-                    words += line[1][0].split()
+            words = result
         logging.debug(f'extracted text for user {mediaitem_user_id} mediaitem {mediaitem_id}: {words}')
-
         if len(words) == 0:
             return None
 

worker/tests/providers/ocr/test_paddle.py

@@ -5,31 +5,42 @@
 from src.providers.ocr.paddlepaddle import PaddleModule
 
 
-@mock.patch('paddleocr.PaddleOCR.__init__', return_value=None)
-def test_paddle_success(_):
+class OCRResult(dict):
+    def __getattr__(self, attr):
+        return self.get(attr, None)
+
+@mock.patch('fastdeploy.vision.ocr.ppocr.PPOCRv3.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.DBDetector.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.Classifier.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.Recognizer.__init__', return_value=None)
+def test_paddle_success(_, __, ___, ____):
     paddle_module = PaddleModule({'det_model_dir':'/det_models','rec_model_dir':'/rec_models','cls_model_dir':'/cls_models'})
     paddle_module.model = mock.MagicMock()
-    paddle_module.model.ocr.return_value = [[[[[106.0, 284.0], [266.0, 284.0], [266.0, 321.0], [106.0, 321.0]], ('offer', 0.99)],
-                                             [[[106.0, 284.0], [266.0, 284.0], [266.0, 321.0], [106.0, 321.0]], ('pizza', 0.96)],
-                                             [[[106.0, 284.0], [266.0, 284.0], [266.0, 321.0], [106.0, 321.0]], ('4.99', 0.94)]]]
+    paddle_module.model.predict.return_value = OCRResult({'text': ['pizza', '4.99', 'offer']})
     result = paddle_module.extract('mediaitem_user_id', 'mediaitem_id', 'photo', 'previews/file_name')
     if 'value' in result.keys():
         sorted(result['value'])
         assert result == {'userId': 'mediaitem_user_id', 'id': 'mediaitem_id', 'name': 'ocr', 'value': ['pizza', '4.99', 'offer']}
 
-@mock.patch('paddleocr.PaddleOCR.__init__', return_value=None)
-def test_paddle_failed_empty_response(_):
+@mock.patch('fastdeploy.vision.ocr.ppocr.PPOCRv3.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.DBDetector.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.Classifier.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.Recognizer.__init__', return_value=None)
+def test_paddle_failed_empty_response(_, __, ___, ____):
     paddle_module = PaddleModule({'det_model_dir':'/det_models','rec_model_dir':'/rec_models','cls_model_dir':'/cls_models'})
     paddle_module.model = mock.MagicMock()
-    paddle_module.model.ocr.return_value = {}
+    paddle_module.model.predict.return_value = None
     result = paddle_module.extract('mediaitem_user_id', 'mediaitem_id', 'photo', 'previews/file_name')
     assert result == None
 
-@mock.patch('paddleocr.PaddleOCR.__init__', return_value=None)
-def test_paddle_failed_exception(_):
+@mock.patch('fastdeploy.vision.ocr.ppocr.PPOCRv3.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.DBDetector.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.Classifier.__init__', return_value=None)
+@mock.patch('fastdeploy.vision.ocr.ppocr.Recognizer.__init__', return_value=None)
+def test_paddle_failed_exception(_, __, ___, ____):
     paddle_module = PaddleModule({'det_model_dir':'/det_models','rec_model_dir':'/rec_models','cls_model_dir':'/cls_models'})
     paddle_module.model = mock.MagicMock()
-    paddle_module.model.ocr.side_effect = Exception('some error')
+    paddle_module.model.predict.side_effect = Exception('some error')
     with pytest.raises(Exception):
         result = paddle_module.extract('mediaitem_user_id', 'mediaitem_id', 'photo', 'previews/file_name')
         assert result == None