first version beta

Author: bigsk05

README.md

@@ -1,3 +1,35 @@
 # arnrs
 
-Aircraft Registration Number Recognition System
+Aircraft Registration Number Recognition System
+
+## Installation
+
+1. Install [Pytorch Engine](https://pytorch.org/get-started/locally/) (CUDA, ROCm or CPU) with anaconda
+
+2. Clone Detectron2 Framework from Github with follow command:
+
+    `git clone https://github.com/facebookresearch/detectron2.git`
+
+3. Run Detectron2 Installtion with follow command:
+
+    `python -m pip install -e detectron2 -i https://pypi.tuna.tsinghua.edu.cn/simple`
+
+4. Install [PaddlePaddle Engine](https://www.paddlepaddle.org.cn/install/quick?docurl=/documentation/docs/zh/install/conda/windows-conda.html) (CUDA, ROCm or CPU) with anaconda
+
+5. Install easyocr package with follow command:
+
+    `pip install easyocr -i https://pypi.tuna.tsinghua.edu.cn/simple`
+
+6. Install paddleocr package with follow command:
+
+    `pip install paddleocr -i https://pypi.tuna.tsinghua.edu.cn/simple`
+
+7. To solve opencv version issue, reinstall opencv and opencv-headless with follow commands (There may be some errors happen, you can ignore it)
+
+    `pip install opencv-python==4.2.0.34 -i https://pypi.tuna.tsinghua.edu.cn/simple`
+
+    `pip install opencv-python-headless==4.2.0.34 -i https://pypi.tuna.tsinghua.edu.cn/simple`
+
+8. Replace packages rely in paddleocr from "tools.infer" to "paddleocr.tools.infer" with text editor, it's a bug from paddleocr
+
+9. Finally enjoy it :)

detector.py

@@ -0,0 +1,39 @@
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.data import MetadataCatalog
+from detectron2.utils.visualizer import ColorMode, Visualizer
+from detectron2 import model_zoo
+
+class Detector:
+    def __init__(self, mode = "zoo", model = "COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml", device = "cpu") -> None:
+        self.cfg = get_cfg()
+
+        if mode == "zoo":
+            # Load model config and pretrained model
+            self.cfg.merge_from_file(model_zoo.get_config_file(model))
+            self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model)
+
+        self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
+        self.cfg.MODEL.DEVICE = device
+
+        self.predictor = DefaultPredictor(self.cfg)
+
+    def image(self, image):
+        predictions = self.predictor(image)
+
+        viz = Visualizer(image[:,:,::-1], metadata = MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]))
+        #instance_mode = ColorMode.IMAGE_BW
+
+        output = viz.draw_instance_predictions(predictions["instances"].to("cpu"))
+
+        result = []
+        for i in range(len(predictions["instances"].scores)):
+            result.append(
+                {
+                    "box": tuple([float(pos) for pos in [obj for obj in predictions["instances"].pred_boxes[i]][0]]),
+                    "score": float(predictions["instances"].scores[i]),
+                    "class": MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]).thing_classes[predictions["instances"].pred_classes[i]]
+                }
+            )
+
+        return output.get_image()[:,:,::-1], tuple(result)

main.py

@@ -0,0 +1,72 @@
+import os
+
+from detector import Detector
+from paddleocr import PaddleOCR
+
+import cv2
+import easyocr
+import requests
+
+class number(object):
+    def __init__(self, gpu=False, times=2):
+        assert type(gpu) is bool
+        assert type(times) is int and times >= 1
+        self.gpu = gpu
+        self.times = times
+
+        # Init detector and OCR
+        self.__detector = Detector(device="gpu" if gpu else "cpu")
+        self.__eocr = easyocr.Reader(['ch_sim', 'en'], gpu=self.gpu)
+        self.__pocr = PaddleOCR(use_angle_cls=True, use_gpu=self.gpu, show_log=False)
+
+    def recognize(self, image):
+        '''
+        Recognize aircraft registration number with detection
+        and ocr powered by pytorch and paddlepaddle engine.
+        :image Accept numpy array or image file path
+        '''
+        if type(image) is str:
+            path = os.path.abspath(image)
+            image = cv2.imread(path)
+
+        result = self.__detector.image(image)
+        # Subjective judgment
+        area_max = 0
+        area_index = 0
+        for i in range(len(result[1])):
+            d = result[1][i]
+            this_area = ((d["box"][1] - d["box"][0]) ** 2 + (d["box"][3] - d["box"][2]) ** 2) ** 0.5
+            if this_area > area_max and result[1][i]["class"] == "airplane":
+                area_max = this_area
+                area_index = i
+
+        i = result[1][area_index]
+        img = image[int(i["box"][1]):int(i["box"][3]), int(i["box"][0]):int(i["box"][2])]
+
+        ocr_result = []
+        ocr_filter = []
+
+        for _ in range(2):
+            eocr_result = self.__eocr.readtext(img, detail=1)
+            pocr_result = self.__pocr.ocr(img, cls=True)
+
+            for e in eocr_result:
+                if e[2] > 0.6 and e[1] not in ocr_filter:
+                    ocr_result.append(
+                        (tuple([tuple(i) for i in e[0]]), e[1], e[2])
+                    )
+                    ocr_filter.append(e[1])
+            for p in pocr_result[0]:
+                if p[1][1] > 0.6 and e[1][0] not in ocr_filter:
+                    ocr_result.append(
+                        (tuple([tuple(i) for i in p[0]]), e[1][0], e[1][1])
+                    )
+                    ocr_filter.append(e[1][0])
+
+        # Read database
+        for i in ocr_result:
+            db = requests.post("http://www.airframes.org/", data={"reg1": i[1]}).text
+            if "No data found on this query." not in db:
+                return i
+
+        return None