From e9113cc3c65d84fbad0b0332ed30dc728d32e939 Mon Sep 17 00:00:00 2001 From: Bigsk Date: Wed, 18 Jan 2023 17:03:53 +0800 Subject: [PATCH] first version beta --- README.md | 34 ++++++++++++++++++++++++- detector.py | 39 +++++++++++++++++++++++++++++ main.py | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 detector.py create mode 100644 main.py diff --git a/README.md b/README.md index afeb28e..d7533a0 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,35 @@ # arnrs -Aircraft Registration Number Recognition System \ No newline at end of file +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 :) \ No newline at end of file diff --git a/detector.py b/detector.py new file mode 100644 index 0000000..882163a --- /dev/null +++ b/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) \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000..d158d22 --- /dev/null +++ b/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 \ No newline at end of file