first commit

.gitignore

@@ -0,0 +1,165 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+.vscode/
+images/
+models/
+logs/

requirements.txt

@@ -0,0 +1,6 @@
+torch 
+torchvision 
+torchaudio
+Pillow
+opencv-python==4.5.1.48
+tqdm

src/apps/char_detection.py

@@ -0,0 +1,123 @@
+'''
+@Author     : Ali Mustofa HALOTEC
+@Module     : Character Detection Faster RCNN
+@Created on : 19 Jul 2022
+'''
+#!/usr/bin/env python3
+# Path: src/apps/char_detection.py
+import os
+import cv2
+import numpy as np
+from PIL import Image
+from src.utils.utils import download_and_unzip_model
+import torch
+import torchvision
+from torchvision import transforms
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+
+
+class CharDetection:
+
+    def __init__(self, root_path:str, model_config:dict) -> None:
+        '''
+        Load model 
+        @params:
+            - root_path:str ->  root of path model
+            - model_config:dict -> config of model {filename, classes, url, file_size}
+        '''
+        self.model_name     = f'{root_path}/{model_config["filename"]}'
+        self.classes        = model_config['classes']
+        self.device         = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+        self.__check_model()
+        self.model          = self.__load_model()
+
+    def __check_model(self, root_path:str, model_config:dict) -> None:
+        if not os.path.isfile(self.model_name):
+            download_and_unzip_model(
+                root_dir    = root_path,
+                name        = model_config['filename'],
+                url         = model_config['url'],
+                file_size   = model_config['file_size'],
+                unzip       = False
+            )
+        else: print('Load model')
+
+    @staticmethod
+    def __image_transform(image) -> torch.Tensor:
+        return transforms.Compose([transforms.ToTensor()])(image)
+
+    def __load_model(self) -> torch.nn.Module:
+        model = self.__fasterrcnn_resnet50_fpn()
+        model.load_state_dict(torch.load(self.model_name, map_location=self.device), False)
+        model.to(self.device)
+        return model.eval()
+
+    def __fasterrcnn_resnet50_fpn(self)-> torch.nn.Module:
+        model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
+        in_features = model.roi_heads.box_predictor.cls_score.in_features
+        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, len(self.classes)+1)
+        return model
+
+    @staticmethod
+    def __filter_threshold(probs:dict, threshold:float) -> dict:
+        num_filtered = (probs['scores']>threshold).float()
+        keep        = (num_filtered == torch.tensor(1)).nonzero().flatten()
+        final_probs = probs
+        final_probs['boxes']    = final_probs['boxes'][keep]
+        final_probs['scores']   = final_probs['scores'][keep]
+        final_probs['labels']   = final_probs['labels'][keep]
+        return final_probs
+
+    @staticmethod
+    def __original_boxes(boxes:torch.Tensor, img_size:tuple,resized:int) -> torch.Tensor:
+        image_width, image_height = img_size[1], img_size[0]
+        boxes = torch.tensor([[
+            (x_min/resized)*image_width, (y_min/resized)*image_height, \
+            (x_max/resized)*image_width, (y_max/resized)*image_height] \
+        for (x_min, y_min, x_max, y_max) in boxes.cpu().numpy()])
+        return boxes
+
+    @staticmethod
+    def __sort_by_boxes(probs:dict) -> dict:
+        x_min_list = [i[0] for i in probs['boxes']]
+        idx = [x_min_list.index(x) for x in sorted(x_min_list)]
+        probs['boxes']    = probs['boxes'][idx]
+        probs['scores']   = probs['scores'][idx]
+        probs['labels']   = probs['labels'][idx]
+        return probs     
+
+    def detect(self, image:np.array, size:int = None, 
+            boxes_ori:bool = False, threshold:float = 0.5, sorted:bool = True) -> dict:
+        '''
+        @params:
+            - image: numpy array of image
+            - size: int of image resize
+            - boxes_ori: bool of original boxes
+            - threshold: float of threshold
+            - sorted: bool of sorted by boxes
+        @return:
+            probs: dict of probs -> {
+                'boxes' : [x_min, y_min, x_max, y_max], 
+                'scores': [float], 
+                'labels': [int]
+            }
+        '''
+        im_shape = (image.shape[0], image.shape[1])
+        image = cv2.resize(image, (size,size)) if size else image
+        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+        image = self.__image_transform(image)
+        with torch.no_grad():
+            probs = self.model([image])[0]
+        probs = self.__filter_threshold(probs, threshold)
+        if boxes_ori and size:
+            probs['boxes'] = self.__original_boxes(probs['boxes'],im_shape, size)
+        if sorted:
+            probs = self.__sort_by_boxes(probs)
+        return {k: v.cpu().numpy() for k, v in probs.items()}
+
+
+if __name__ == '__main__':
+    char_detection = CharDetection('./models/text_detection.ali', ['text'])
+    image  = cv2.imread('./images/1.jpg')
+    results = char_detection.detect(image, size=244, boxes_ori=True, threshold=0.01)
+    print(results)