├── .flake8 ├── .github └── workflows │ └── build_packages.yml ├── .gitignore ├── .isort.cfg ├── .pre-commit-config.yaml ├── .prettierrc ├── README.md ├── modelplace_api ├── __init__.py ├── _rle_mask.py ├── base_classes.py ├── colors.py ├── objects.py ├── text_styles │ ├── Montserrat-Bold.ttf │ ├── Montserrat-Medium.ttf │ └── Montserrat-Regular.ttf ├── utils.py └── visualization.py ├── pyproject.toml ├── requirements-dev.txt ├── setup.py └── tests ├── decoded_binary_mask.json ├── test_rle.py └── test_rle_backward_compatibility.py /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | ignore = E203, E266, E501, W503 3 | # line length is intentionally set to 80 here because black uses Bugbear 4 | # See https://github.com/psf/black/blob/master/README.md#line-length for more details 5 | max-line-length = 88 6 | max-complexity = 18 7 | select = B,C,E,F,W,T4,B9 8 | # We need to configure the mypy.ini because the flake8-mypy's default 9 | # options don't properly override it, so if we don't specify it we get 10 | # half of the config from mypy.ini and half from flake8-mypy. 11 | mypy_config = mypy.ini 12 | -------------------------------------------------------------------------------- /.github/workflows/build_packages.yml: -------------------------------------------------------------------------------- 1 | name: Build wheels for modelplace-api on Windows and Unix 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | branches: 9 | - main 10 | 11 | jobs: 12 | build: 13 | runs-on: ${{ matrix.os }} 14 | 15 | strategy: 16 | fail-fast: false 17 | matrix: 18 | os: [windows-latest, ubuntu-latest, macos-latest] 19 | python-version: [3.6, 3.7, 3.8, 3.9] 20 | platform: [x64] 21 | 22 | steps: 23 | - name: Checkout 24 | uses: actions/checkout@v2 25 | with: 26 | fetch-depth: 0 27 | 28 | - name: Set up Python ${{ matrix.python-version }} 29 | uses: actions/setup-python@v2 30 | with: 31 | python-version: ${{ matrix.python-version }} 32 | architecture: ${{ matrix.platform }} 33 | 34 | - name: Build the wheel on Windows 35 | if: ${{ matrix.os == 'windows-latest' }} 36 | run: | 37 | python --version 38 | cd ${{ github.workspace }} 39 | pip3 install wheel pytest && python setup.py bdist_wheel 40 | shell: cmd 41 | 42 | - name: Build the wheel on Unix 43 | if: ${{ matrix.os != 'windows-latest' && matrix.python-version != '3.6' }} 44 | run: | 45 | python --version 46 | cd ${{ github.workspace }} 47 | pip3 install wheel pytest && python setup.py bdist_wheel 48 | 49 | - name: Build the wheel on Unix using python3.6 50 | if: ${{ matrix.os != 'windows-latest' && matrix.python-version == '3.6' }} 51 | run: | 52 | python --version 53 | cd ${{ github.workspace }} 54 | pip3 install wheel pytest numpy==1.19.3 && python setup.py bdist_wheel 55 | 56 | - name: Install the package on Windows 57 | if: ${{ matrix.os == 'windows-latest' }} 58 | run: | 59 | cd ${{ github.workspace }}/dist 60 | FOR %%i in (*.whl) DO python -m pip install %%i[vis-windows] 61 | shell: cmd 62 | 63 | - name: Install the package on Unix 64 | if: ${{ matrix.os != 'windows-latest' }} 65 | run: for i in ${{ github.workspace }}/dist/*.whl; do python -m pip install ${i}[vis]; done 66 | 67 | - name: Test the package on Windows 68 | if: ${{ matrix.os == 'windows-latest' }} 69 | run: cd ${{ github.workspace }}/tests && python -m pytest test_rle.py -s 70 | shell: cmd 71 | 72 | - name: Test the package on Unix 73 | if: ${{ matrix.os != 'windows-latest' }} 74 | run: cd ${{ github.workspace }}/tests && python -m pytest test_rle_backward_compatibility.py -s 75 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *.py[cod] 3 | *$py.class 4 | .idea/ 5 | *.so 6 | *.o 7 | dist 8 | build 9 | *.egg-info 10 | *.orig 11 | .vscode 12 | .nox 13 | .mypy_cache 14 | *.code-workspace 15 | -------------------------------------------------------------------------------- /.isort.cfg: -------------------------------------------------------------------------------- 1 | [settings] 2 | indent=4 3 | known_third_party=cv2,loguru,numpy,pycocotools,pydantic,setuptools 4 | line_length=88 5 | use_parentheses=True 6 | multi_line_output=3 7 | combine_as_imports=False 8 | force_grid_wrap=0 9 | combine_star=True 10 | use_parentheses=True 11 | from_first=False 12 | include_trailing_comma=True 13 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | default_language_version: 2 | python: python3.7 3 | repos: 4 | - repo: https://github.com/pre-commit/pre-commit-hooks 5 | rev: v2.5.0 6 | hooks: 7 | - id: check-yaml 8 | - id: trailing-whitespace 9 | files: '\.(py|md|txt)$' 10 | - id: end-of-file-fixer 11 | files: '\.(py|md|txt)$' 12 | - id: check-ast 13 | - id: mixed-line-ending 14 | - id: requirements-txt-fixer 15 | 16 | - repo: https://github.com/Lucas-C/pre-commit-hooks 17 | rev: v1.1.7 18 | hooks: 19 | - id: remove-tabs 20 | files: \.(py,md)$ 21 | - repo: https://github.com/asottile/add-trailing-comma 22 | rev: v2.0.1 23 | hooks: 24 | - id: add-trailing-comma 25 | files: \.py$ 26 | - repo: https://github.com/asottile/seed-isort-config 27 | rev: v2.1.1 28 | hooks: 29 | - id: seed-isort-config 30 | files: \.py$ 31 | - repo: local 32 | hooks: 33 | - id: isort 34 | name: isort 35 | entry: python3 -m isort -y --settings-path pyproject.toml 36 | types: [python] 37 | language: system 38 | - repo: local 39 | hooks: 40 | - id: prettier 41 | name: prettier 42 | language: docker_image 43 | entry: tmknom/prettier:2.0.5 --parser=markdown --write --config ./.prettierrc 44 | files: (?!(styles|ci)/)\.md$ 45 | - repo: local 46 | hooks: 47 | - id: prettier 48 | name: prettier 49 | language: docker_image 50 | entry: tmknom/prettier:2.0.5 --parser=yaml --write --config ./.prettierrc 51 | files: \.ya?ml$ 52 | - repo: local 53 | hooks: 54 | - id: black 55 | name: black 56 | types: [python] 57 | language: system 58 | entry: black --config pyproject.toml 59 | -------------------------------------------------------------------------------- /.prettierrc: -------------------------------------------------------------------------------- 1 | trailingComma: "es5" 2 | tabWidth: 2 3 | semi: false 4 | singleQuote: false 5 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Modelplace Api 2 | 3 | This repository contains the required API and some useful functions for launching [samples for OAK](https://github.com/opencv-ai/oak-model-samples) 4 | 5 | ## Usage 6 | 7 | It's a dependency package which is built automatically when `setup.py` for OAK sample is called. 8 | -------------------------------------------------------------------------------- /modelplace_api/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.4.20" 2 | 3 | import modelplace_api.utils 4 | import modelplace_api.visualization 5 | from modelplace_api.base_classes import * 6 | from modelplace_api.objects import * 7 | -------------------------------------------------------------------------------- /modelplace_api/_rle_mask.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def encode_binary_mask(binary_mask: np.ndarray) -> dict: 5 | h, w = binary_mask.shape 6 | binary_mask = binary_mask.flatten(order="F") 7 | mask_len = len(binary_mask) 8 | counts_list = [] 9 | pos = 0 10 | 11 | # RLE encoding 12 | counts_list.append(1) 13 | diffs = np.logical_xor(binary_mask[0 : mask_len - 1], binary_mask[1:mask_len]) 14 | for diff in diffs: 15 | if diff: 16 | pos += 1 17 | counts_list.append(1) 18 | else: 19 | counts_list[pos] += 1 20 | 21 | # Encoding an array into a byte array 22 | # if array starts from 1. start with 0 counts for 0 23 | if binary_mask[0] == 1: 24 | counts_list = [0] + counts_list 25 | counts = [] 26 | more = True 27 | i = 0 28 | for i in range(len(counts_list)): 29 | x = counts_list[i] 30 | if i > 2: 31 | x -= counts_list[i - 2] 32 | more = True 33 | while more: 34 | c = x & 0x1F 35 | x >>= 5 36 | more = x != -1 if (c & 0x10) else x != 0 37 | if more: 38 | c |= 0x20 39 | c += 48 # shift for byte 40 | counts.append(chr(c)) 41 | 42 | return { 43 | "size": [h, w], 44 | "counts": "".join(counts).encode("utf-8"), 45 | } 46 | 47 | 48 | def decode_coco_rle(rle_mask: dict) -> np.ndarray: 49 | binaries = rle_mask["counts"] 50 | resulted_mask = [] 51 | more = True 52 | i = 0 53 | k = 0 54 | x = 0 55 | prev_x = 0 56 | prev_prev_x = 0 57 | while k < len(binaries): 58 | prev_prev_x = prev_x 59 | prev_x = x 60 | x = 0 61 | m = 0 62 | more = True 63 | while more: 64 | c = binaries[k] - 48 # shift for byte 65 | x |= (c & 0x1F) << 5 * m # number of values (0 or q) 66 | more = c & 0x20 67 | k += 1 68 | m += 1 69 | if not more and (c & 0x10): 70 | x |= -1 << 5 * m 71 | if i > 2: 72 | x += prev_prev_x # add previous counts 73 | i += 1 74 | value = not i % 2 # alternation of 0 and 1 75 | resulted_mask += [value] * x # create a sequence of values 76 | return np.array(resulted_mask, dtype=np.uint8).reshape(rle_mask["size"], order="F") 77 | -------------------------------------------------------------------------------- /modelplace_api/base_classes.py: -------------------------------------------------------------------------------- 1 | import traceback 2 | from abc import ABC, abstractmethod 3 | from typing import Any, Tuple 4 | 5 | from loguru import logger 6 | 7 | from .objects import Device 8 | 9 | 10 | class BaseModel(ABC): 11 | def __init__( 12 | self, 13 | model_path: str, 14 | model_name: str = "", 15 | model_description: str = "", 16 | **kwargs, 17 | ): 18 | self.model_path = model_path 19 | self.model_name = model_name 20 | self.model_description = model_description 21 | self.model = None 22 | 23 | @logger.catch(onerror=lambda _: traceback.print_exc()) 24 | @abstractmethod 25 | def preprocess(self, *args, **kwargs) -> Tuple[Any]: 26 | raise NotImplementedError 27 | 28 | @logger.catch(onerror=lambda _: traceback.print_exc()) 29 | @abstractmethod 30 | def postprocess(self, *args, **kwargs) -> Tuple[Any]: 31 | raise NotImplementedError 32 | 33 | @logger.catch(onerror=lambda _: traceback.print_exc()) 34 | @abstractmethod 35 | def model_load(self, device: Device) -> None: 36 | raise NotImplementedError 37 | 38 | def forward(self, data: Any) -> Any: 39 | result = self.model(data) 40 | return result 41 | 42 | @logger.catch(onerror=lambda _: traceback.print_exc()) 43 | def __call__(self, *args, **kwargs) -> Any: 44 | return self.forward(*args) 45 | 46 | @logger.catch(onerror=lambda _: traceback.print_exc()) 47 | @abstractmethod 48 | def process_sample(self, image: Any) -> Any: 49 | raise NotImplementedError 50 | -------------------------------------------------------------------------------- /modelplace_api/colors.py: -------------------------------------------------------------------------------- 1 | RGBA_COLORS = [ 2 | (255, 74, 70, 1), 3 | (0, 137, 65, 1), 4 | (0, 111, 166, 1), 5 | (163, 0, 89, 1), 6 | (0, 0, 166, 1), 7 | (0, 77, 67, 1), 8 | (90, 0, 7, 1), 9 | (27, 68, 0, 1), 10 | (59, 93, 255, 1), 11 | (74, 59, 83, 1), 12 | (186, 9, 0, 1), 13 | (107, 121, 0, 1), 14 | (185, 3, 170, 1), 15 | (209, 97, 0, 1), 16 | (48, 0, 24, 1), 17 | (10, 166, 216, 1), 18 | (1, 51, 73, 1), 19 | (0, 132, 111, 1), 20 | (160, 121, 191, 1), 21 | (204, 7, 68, 1), 22 | (0, 30, 9, 1), 23 | (0, 72, 156, 1), 24 | (111, 0, 98, 1), 25 | (69, 109, 117, 1), 26 | (183, 123, 104, 1), 27 | (122, 135, 161, 1), 28 | (120, 141, 102, 1), 29 | (136, 85, 120, 1), 30 | (209, 87, 160, 1), 31 | (69, 102, 72, 1), 32 | (136, 111, 76, 1), 33 | (52, 54, 45, 1), 34 | (99, 99, 117, 1), 35 | (87, 83, 41, 1), 36 | (176, 91, 111, 1), 37 | (59, 151, 0, 1), 38 | (30, 110, 0, 1), 39 | (121, 0, 215, 1), 40 | (167, 117, 0, 1), 41 | (99, 103, 169, 1), 42 | (160, 88, 55, 1), 43 | (107, 0, 44, 1), 44 | (119, 38, 0, 1), 45 | (84, 158, 121, 1), 46 | (114, 65, 143, 1), 47 | (153, 173, 192, 1), 48 | (58, 36, 101, 1), 49 | (146, 35, 41, 1), 50 | (64, 78, 85, 1), 51 | (203, 126, 152, 1), 52 | (50, 78, 114, 1), 53 | (106, 58, 76, 1), 54 | (131, 171, 88, 1), 55 | (0, 75, 40, 1), 56 | (128, 108, 102, 1), 57 | (191, 86, 80, 1), 58 | (232, 48, 0, 1), 59 | (102, 121, 109, 1), 60 | (91, 78, 81, 1), 61 | (255, 104, 50, 1), 62 | (1, 44, 88, 1), 63 | (122, 123, 255, 1), 64 | (214, 142, 1, 1), 65 | (53, 51, 57, 1), 66 | (120, 175, 161, 1), 67 | (131, 115, 147, 1), 68 | (148, 58, 77, 1), 69 | (149, 86, 189, 1), 70 | (106, 113, 74, 1), 71 | (2, 82, 95, 1), 72 | (233, 129, 118, 1), 73 | (61, 79, 68, 1), 74 | (126, 100, 5, 1), 75 | (2, 104, 78, 1), 76 | (150, 43, 117, 1), 77 | (216, 106, 120, 1), 78 | (62, 137, 190, 1), 79 | (202, 131, 78, 1), 80 | (81, 138, 135, 1), 81 | (91, 17, 60, 1), 82 | (85, 129, 59, 1), 83 | (0, 0, 95, 1), 84 | (169, 115, 153, 1), 85 | (75, 129, 96, 1), 86 | (89, 115, 138, 1), 87 | (100, 49, 39, 1), 88 | (107, 148, 170, 1), 89 | (81, 160, 88, 1), 90 | (164, 91, 2, 1), 91 | (76, 96, 1, 1), 92 | (156, 105, 102, 1), 93 | (100, 84, 123, 1), 94 | (151, 151, 158, 1), 95 | (0, 106, 102, 1), 96 | (57, 20, 6, 1), 97 | (0, 69, 210, 1), 98 | (0, 108, 49, 1), 99 | (124, 101, 113, 1), 100 | (32, 59, 60, 1), 101 | (103, 17, 144, 1), 102 | (107, 58, 100, 1), 103 | (55, 69, 39, 1), 104 | (200, 98, 64, 1), 105 | (41, 96, 124, 1), 106 | (125, 90, 68, 1), 107 | (184, 129, 131, 1), 108 | (170, 81, 153, 1), 109 | (167, 69, 113, 1), 110 | (120, 158, 201, 1), 111 | (109, 128, 186, 1), 112 | (149, 63, 0, 1), 113 | (0, 49, 9, 1), 114 | (0, 96, 205, 1), 115 | (137, 85, 99, 1), 116 | (91, 50, 19, 1), 117 | (167, 111, 66, 1), 118 | (137, 65, 46, 1), 119 | (26, 58, 42, 1), 120 | (73, 75, 90, 1), 121 | (168, 140, 133, 1), 122 | (190, 71, 0, 1), 123 | (101, 129, 136, 1), 124 | (131, 164, 133, 1), 125 | (69, 60, 35, 1), 126 | (71, 103, 93, 1), 127 | (58, 63, 0, 1), 128 | ] 129 | -------------------------------------------------------------------------------- /modelplace_api/objects.py: -------------------------------------------------------------------------------- 1 | import enum 2 | from typing import List 3 | 4 | import pydantic 5 | 6 | 7 | class Point(pydantic.BaseModel): 8 | x: int 9 | y: int 10 | 11 | 12 | class TextPolygon(pydantic.BaseModel): 13 | points: List[Point] 14 | text: str = "" 15 | 16 | 17 | class BBox(pydantic.BaseModel): 18 | x1: float 19 | y1: float 20 | x2: float 21 | y2: float 22 | score: float 23 | class_name: str 24 | 25 | 26 | class TrackBBox(pydantic.BaseModel): 27 | x1: float 28 | y1: float 29 | x2: float 30 | y2: float 31 | score: float 32 | class_name: str 33 | track_number: int 34 | 35 | 36 | class COCOBBox(pydantic.BaseModel): 37 | x1: float 38 | y1: float 39 | x2: float 40 | y2: float 41 | score: float 42 | class_name: str 43 | area: float 44 | is_crowd: int 45 | 46 | 47 | class Mask(pydantic.BaseModel): 48 | mask: dict 49 | classes: list 50 | 51 | 52 | class InstanceMask(pydantic.BaseModel): 53 | detections: List[BBox] 54 | masks: List[Mask] 55 | 56 | 57 | class Joint(pydantic.BaseModel): 58 | x: int 59 | y: int 60 | class_name: str 61 | score: float 62 | 63 | 64 | class Link(pydantic.BaseModel): 65 | joint_a: Joint 66 | joint_b: Joint 67 | 68 | 69 | class Pose(pydantic.BaseModel): 70 | score: float 71 | links: List[Link] 72 | skeleton_parts: List 73 | 74 | 75 | class COCOPose(pydantic.BaseModel): 76 | score: float 77 | links: List[Link] 78 | skeleton_parts: List 79 | # GT PARTS 80 | bbox: COCOBBox 81 | 82 | 83 | class COCOInstanceMask(pydantic.BaseModel): 84 | detections: List[COCOBBox] 85 | masks: List[Mask] 86 | 87 | 88 | class Label(pydantic.BaseModel): 89 | score: float 90 | class_name: str 91 | 92 | 93 | class AgeGenderLabel(pydantic.BaseModel): 94 | bbox: BBox 95 | age: int 96 | genders: List[Label] 97 | 98 | 99 | class EmotionLabel(pydantic.BaseModel): 100 | bbox: BBox 101 | emotions: List[Label] 102 | 103 | 104 | class Landmarks(pydantic.BaseModel): 105 | bbox: BBox 106 | keypoints: List[Point] 107 | 108 | 109 | class VideoFrame(pydantic.BaseModel): 110 | number: int 111 | boxes: List[TrackBBox] 112 | 113 | 114 | class CountableVideoFrame(VideoFrame): 115 | people_in: int 116 | people_out: int 117 | 118 | 119 | class File(pydantic.BaseModel): 120 | data: bytes 121 | extension: str 122 | 123 | @pydantic.validator("extension") 124 | def validate_extension(cls, v) -> str: 125 | return v.split(".")[-1] 126 | 127 | 128 | class Device(enum.Enum): 129 | gpu = enum.auto() 130 | cpu = enum.auto() 131 | 132 | 133 | @enum.unique 134 | class TaskType(enum.Enum): 135 | detection = enum.auto() 136 | segmentation = enum.auto() 137 | pose_estimation = enum.auto() 138 | tracking = enum.auto() 139 | text_detection = enum.auto() 140 | unknown = enum.auto() 141 | people_counting = enum.auto() 142 | classification = enum.auto() 143 | landmark_detection = enum.auto() 144 | age_gender_recognition = enum.auto() 145 | emotion_recognition = enum.auto() 146 | mesh_detection = enum.auto() 147 | background_removal = enum.auto() 148 | instance_segmentation = enum.auto() 149 | image_processing = enum.auto() 150 | -------------------------------------------------------------------------------- /modelplace_api/text_styles/Montserrat-Bold.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/opencv-ai/modelplace-api/8bbf7c6b97b0b63cee9d8a264de2f103afc8a8ca/modelplace_api/text_styles/Montserrat-Bold.ttf -------------------------------------------------------------------------------- /modelplace_api/text_styles/Montserrat-Medium.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/opencv-ai/modelplace-api/8bbf7c6b97b0b63cee9d8a264de2f103afc8a8ca/modelplace_api/text_styles/Montserrat-Medium.ttf -------------------------------------------------------------------------------- /modelplace_api/text_styles/Montserrat-Regular.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/opencv-ai/modelplace-api/8bbf7c6b97b0b63cee9d8a264de2f103afc8a8ca/modelplace_api/text_styles/Montserrat-Regular.ttf -------------------------------------------------------------------------------- /modelplace_api/utils.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | import numpy as np 4 | from loguru import logger 5 | 6 | try: 7 | from pycocotools import mask 8 | 9 | encode_binary_mask = mask.encode 10 | decode_coco_rle = mask.decode 11 | except ImportError: 12 | logger.warning( 13 | "The 'pycocotools' package wasn't found. Slow encoding and decoding are used for the RLE mask.", 14 | ) 15 | 16 | from ._rle_mask import encode_binary_mask, decode_coco_rle 17 | 18 | 19 | def is_numpy_array_equal(result: np.ndarray, gt: np.ndarray, error: float) -> bool: 20 | """ 21 | Applies element-wise comparison of two ndarrays and defines 22 | whether the ratio of matched elements is eligible 23 | """ 24 | if result.shape != gt.shape: 25 | raise RuntimeError( 26 | f'"result" and "gt" shapes are different ({result.shape} vs {gt.shape}) - must be the same', 27 | ) 28 | 29 | if np.issubdtype(result.dtype, np.integer): 30 | matched = np.equal(result, gt).sum() 31 | else: 32 | matched = np.isclose(result, gt, rtol=error).sum() 33 | return matched / result.size >= 1 - error 34 | 35 | 36 | def is_equal(result: Any, gt: Any, error: float = 0.001) -> bool: 37 | if type(result) != type(gt): 38 | raise TypeError 39 | 40 | ret = True 41 | if isinstance(result, dict): 42 | for key in result: 43 | ret = ret and is_equal(result[key], gt[key], error) 44 | elif isinstance(result, list): 45 | if not result: 46 | ret = ret and not bool(len(gt)) 47 | else: 48 | for r, g in zip(result, gt): 49 | ret = ret and is_equal(r, g, error) 50 | elif isinstance(result, str): 51 | ret = ret and result == gt 52 | elif isinstance(result, bytes): 53 | result = result.decode("utf-8") 54 | gt = gt.decode("utf-8") 55 | ret = ret and result == gt 56 | elif isinstance(result, np.ndarray): 57 | ret = ret and is_numpy_array_equal(result, gt, error) 58 | else: 59 | ret = ret and np.isclose(result, gt, rtol=error) 60 | return ret 61 | 62 | 63 | def prepare_mask(result_mask: np.ndarray) -> dict: 64 | masks = { 65 | "binary": [], 66 | "classes": [], 67 | } 68 | for unique in np.unique(result_mask): 69 | binary_mask = np.zeros(shape=result_mask.shape, dtype=np.uint8) 70 | binary_mask[result_mask == unique] = 1 71 | masks["binary"].append(encode_binary_mask(np.asfortranarray(binary_mask))) 72 | masks["classes"].append(int(unique)) 73 | return masks 74 | -------------------------------------------------------------------------------- /modelplace_api/visualization.py: -------------------------------------------------------------------------------- 1 | import os 2 | from collections import defaultdict 3 | from typing import Generator, List 4 | 5 | import numpy as np 6 | from loguru import logger 7 | 8 | from .colors import RGBA_COLORS 9 | from .objects import ( 10 | AgeGenderLabel, 11 | BBox, 12 | CountableVideoFrame, 13 | EmotionLabel, 14 | InstanceMask, 15 | Label, 16 | Landmarks, 17 | Mask, 18 | Point, 19 | Pose, 20 | TextPolygon, 21 | VideoFrame, 22 | ) 23 | from .utils import decode_coco_rle 24 | 25 | try: 26 | import cv2 27 | import imageio 28 | import skvideo 29 | from PIL import Image, ImageDraw, ImageFont 30 | except ImportError: 31 | logger.warning( 32 | "Some dependencies is invalid. " 33 | "Please install this package with extra requiements. " 34 | "For unix: pip install modelplace-api[vis] " 35 | "For windows: pip install modelplace-api[vis-windows]", 36 | ) 37 | 38 | 39 | FFMPEG_OUTPUT_DICT = { 40 | "-vcodec": "libx264", 41 | "-vf": "format=yuv420p", 42 | "-movflags": "+faststart", 43 | } 44 | text_style_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)), "text_styles") 45 | 46 | MONTSERATT_BOLD_TTF_PATH = os.path.join(text_style_dir, "Montserrat-Bold.ttf") 47 | MONTSERATT_REGULAR_TTF_PATH = os.path.join(text_style_dir, "Montserrat-Regular.ttf") 48 | MONTSERATT_MEDIUM_TTF_PATH = os.path.join(text_style_dir, "Montserrat-Medium.ttf") 49 | 50 | BACKGROUND_COLOR = (79, 79, 79, 1) 51 | CORALL_COLOR = (236, 113, 108, 1) 52 | DELIMITER_COLOR = (255, 255, 255, 1) 53 | WHITE_TEXT_COLOR = (255, 255, 255, 1) 54 | BLACK_TEXT_COLOR = (0, 0, 0, 1) 55 | LAVANDER_COLOR = (125, 211, 210, 1) 56 | DARK_PINK_COLOR = (109, 84, 199, 1) 57 | NORM_HEIGHT = 591 58 | CLASS_BOX_WIDTH = 80 59 | CLASS_BOX_HEIGHT = 40 60 | INFO_BOX_WIDTH = 60 61 | INFO_BOX_HEIGHT = 30 62 | PROB_BOX_HEIGHT = PROB_BOX_WIDTH = 70 63 | TEXT_OFFSET_X = 16 64 | TEXT_OFFSET_Y = 8 65 | BOX_CORNER_OFFSET = 24 66 | CLASS_BOX_CORNER_OFFSET = 13 67 | CLASS_BOX_TEXT_OFFSET_Y = 12 68 | CLASS_BOX_TEXT_OFFSET_X = 18 69 | COMMENT_TEXT_SIZE = 16 70 | PROB_DIGIT_TEXT_SIZE = 22 71 | PROB_DIGIT_OFFSET_Y = 14 72 | PROB_COMMENT_OFFSET_Y = 8 73 | PROB_COMMENT_OFFSET_X = 10 74 | INFO_TEXT_SIZE = 12 75 | CLASS_TEXT_SIZE = 14 76 | LINE_THINKNESS = 3 77 | POSE_EDGE_THINKNESS = 6 78 | POSE_POINT_SIZE = 10 79 | MESH_POINT_SIZE = 1.5 80 | 81 | 82 | def get_text_width(image, text, text_size, ttf_path): 83 | return ImageDraw.Draw(Image.fromarray(image)).textsize( 84 | text, font=ImageFont.truetype(ttf_path, text_size), 85 | )[0] 86 | 87 | 88 | def add_text( 89 | image: np.ndarray, 90 | text: str, 91 | coords: list, 92 | text_size: int, 93 | text_color: tuple, 94 | ttf_path: str, 95 | box_w: int, 96 | ) -> np.ndarray: 97 | pil_img = Image.fromarray(image) 98 | draw = ImageDraw.Draw(pil_img) 99 | montserrat = ImageFont.truetype(ttf_path, text_size) 100 | text_w = get_text_width(image, text, text_size, ttf_path) 101 | coords[0] = coords[0] + int((box_w - text_w) / 2) 102 | draw.text(coords, text, font=montserrat, fill=text_color) 103 | return np.array(pil_img) 104 | 105 | 106 | def add_probability_box( 107 | image: np.ndarray, 108 | probability: str, 109 | text: str, 110 | coords: List, 111 | delimiter: bool = False, 112 | box_color: tuple = BACKGROUND_COLOR, 113 | ) -> np.ndarray: 114 | img_h, img_w, _ = image.shape 115 | scale = min([img_w, img_h]) / NORM_HEIGHT 116 | prob_size = int(scale * PROB_DIGIT_TEXT_SIZE) 117 | text_size = int(scale * COMMENT_TEXT_SIZE) 118 | box_x1, box_y1, box_x2, box_y2 = coords 119 | box_w = box_x2 - box_x1 120 | box_h = box_y2 - box_y1 121 | text_offset_y1 = int(scale * PROB_DIGIT_OFFSET_Y) 122 | text_offset_y2 = int(box_h - scale * (COMMENT_TEXT_SIZE + 10)) 123 | 124 | image = cv2.rectangle( 125 | image, 126 | (int(box_x1), int(box_y1)), 127 | (int(box_x2), int(box_y2)), 128 | box_color, 129 | thickness=-1, 130 | ) 131 | image = add_text( 132 | image, 133 | probability, 134 | [box_x1, box_y1 + text_offset_y1], 135 | prob_size, 136 | WHITE_TEXT_COLOR, 137 | MONTSERATT_BOLD_TTF_PATH, 138 | box_w, 139 | ) 140 | image = add_text( 141 | image, 142 | text, 143 | [box_x1, box_y1 + text_offset_y2], 144 | text_size, 145 | WHITE_TEXT_COLOR, 146 | MONTSERATT_REGULAR_TTF_PATH, 147 | box_w, 148 | ) 149 | if delimiter: 150 | line_offset_x = int(scale * TEXT_OFFSET_X) 151 | image = cv2.line( 152 | image, 153 | (box_x1 + line_offset_x, box_y2), 154 | (box_x2 - line_offset_x, box_y2), 155 | DELIMITER_COLOR, 156 | 1, 157 | cv2.LINE_AA, 158 | ) 159 | return image 160 | 161 | 162 | def add_class_box( 163 | image: np.ndarray, 164 | background_color: tuple, 165 | text: str, 166 | box_w: int, 167 | box_number: int = 0, 168 | ) -> np.ndarray: 169 | img_h, img_w, _ = image.shape 170 | scale = min([img_w, img_h]) / NORM_HEIGHT 171 | text_size = int(scale * CLASS_TEXT_SIZE) 172 | box_offset_x = box_offset_y = int(scale * CLASS_BOX_CORNER_OFFSET) 173 | box_h = int(scale * CLASS_BOX_HEIGHT) 174 | text_offset_y = int(scale * CLASS_BOX_TEXT_OFFSET_Y) 175 | 176 | box_x2 = img_w - box_offset_x 177 | box_x1 = box_x2 - box_w 178 | box_y2 = box_offset_y + (1 + box_number) * box_h 179 | box_y1 = box_y2 - box_h 180 | image = cv2.rectangle( 181 | image, 182 | (int(box_x1), int(box_y1)), 183 | (int(box_x2), int(box_y2 - 2)), 184 | background_color, 185 | thickness=-1, 186 | ) 187 | image = add_text( 188 | image, 189 | text, 190 | [box_x1, box_y1 + text_offset_y], 191 | text_size, 192 | WHITE_TEXT_COLOR, 193 | MONTSERATT_MEDIUM_TTF_PATH, 194 | box_w, 195 | ) 196 | return image 197 | 198 | 199 | def add_bbox(image: np.ndarray, coords: List, box_color: tuple) -> np.ndarray: 200 | (x1, y1, x2, y2) = coords 201 | img_h, img_w, _ = image.shape 202 | scale = min([img_w, img_h]) / NORM_HEIGHT 203 | thickness = int(scale * LINE_THINKNESS) 204 | image = cv2.rectangle( 205 | image, 206 | (int(x1), int(y1)), 207 | (int(x2), int(y2)), 208 | box_color, 209 | thickness=thickness, 210 | lineType=cv2.LINE_AA, 211 | ) 212 | return image 213 | 214 | 215 | def add_info( 216 | image: np.ndarray, coords: List, box_color: tuple, text: str, text_color: tuple, 217 | ) -> np.ndarray: 218 | img_h, img_w, _ = image.shape 219 | scale = min([img_w, img_h]) / NORM_HEIGHT 220 | coords = coords[0] - int(2 * scale), coords[1] + int(2 * scale) 221 | text_size = int(scale * INFO_TEXT_SIZE) 222 | box_w = int(scale * INFO_BOX_WIDTH) 223 | box_h = int(scale * INFO_BOX_HEIGHT) 224 | text_offset_y = int(scale * TEXT_OFFSET_Y) 225 | text_w = get_text_width(image, text, text_size, MONTSERATT_BOLD_TTF_PATH) 226 | box_w = max(box_w, text_w + text_offset_y * 2) 227 | image = cv2.rectangle( 228 | image, 229 | (int(coords[0]), int(coords[1] - box_h)), 230 | (int(coords[0] + box_w), int(coords[1])), 231 | box_color, 232 | thickness=-1, 233 | ) 234 | image = add_text( 235 | image, 236 | text, 237 | [coords[0], coords[1] - box_h + text_offset_y], 238 | text_size, 239 | text_color, 240 | MONTSERATT_BOLD_TTF_PATH, 241 | box_w, 242 | ) 243 | return image 244 | 245 | 246 | def add_pose(image: np.ndarray, detection: Pose, edge_color: tuple) -> np.ndarray: 247 | img_h, img_w, _ = image.shape 248 | scale = min([img_w, img_h]) / NORM_HEIGHT 249 | edge_thinkness = int(scale * POSE_EDGE_THINKNESS) 250 | for link in detection.links: 251 | if ( 252 | link.joint_b.x == link.joint_b.y == 0 253 | or link.joint_a.x == link.joint_a.y == 0 254 | ): 255 | continue 256 | 257 | cv2.line( 258 | image, 259 | (int(link.joint_a.x), int(link.joint_a.y)), 260 | (int(link.joint_b.x), int(link.joint_b.y)), 261 | edge_color, 262 | edge_thinkness, 263 | ) 264 | 265 | joints = [ 266 | joint for link in detection.links for joint in [link.joint_a, link.joint_b] 267 | ] 268 | unique_joints = [ 269 | a 270 | for i, a in enumerate(joints) 271 | if not any(a.class_name == b.class_name for b in joints[:i]) 272 | ] 273 | for joint in unique_joints: 274 | if joint.x == joint.y == 0: 275 | continue 276 | image = add_point(image, joint, LAVANDER_COLOR, POSE_POINT_SIZE) 277 | return image 278 | 279 | 280 | def add_point( 281 | image: np.ndarray, point: Point, color: tuple, point_size: int = 4, 282 | ) -> np.ndarray: 283 | img_h, img_w, _ = image.shape 284 | scale = min([img_w, img_h]) / NORM_HEIGHT 285 | cv2.circle( 286 | image, (int(point.x), int(point.y)), int(scale * point_size), color, -1, 287 | ) 288 | return image 289 | 290 | 291 | def add_poly( 292 | image: np.ndarray, poly: TextPolygon, color: tuple, thinkness: int, 293 | ) -> np.ndarray: 294 | img_h, img_w, _ = image.shape 295 | scale = min([img_w, img_h]) / NORM_HEIGHT 296 | thinkness = int(scale * thinkness) 297 | points = np.array([[point.x, point.y] for point in poly.points], dtype=np.int32) 298 | image = cv2.polylines(image, [points], True, color, thinkness) 299 | return image 300 | 301 | 302 | def add_mask(image: np.ndarray, idx: np.ndarray, color: tuple) -> np.ndarray: 303 | mask = np.zeros_like(image).astype(np.uint8) 304 | mask[idx] = color[: mask.shape[2]] 305 | alpha = 0.8 306 | beta = 0.4 307 | image = alpha * image + beta * mask 308 | image = image * (2 - np.max(image) / 255) 309 | return image.astype(np.uint8) 310 | 311 | 312 | def add_instance_mask(image: np.ndarray, mask: np.ndarray): 313 | alpha = 0.5 314 | beta = 0.5 315 | image = alpha * image + beta * mask 316 | return image.astype(np.uint8) 317 | 318 | 319 | def add_legend( 320 | image: np.ndarray, classes: List, picked_color: tuple, picked_class_number: int, 321 | ) -> np.ndarray: 322 | 323 | img_h, img_w, _ = image.shape 324 | scale = min([img_w, img_h]) / NORM_HEIGHT 325 | text_size = int(scale * CLASS_TEXT_SIZE) 326 | text_offset_x = int(scale * CLASS_BOX_TEXT_OFFSET_X) 327 | box_w = max( 328 | [int(scale * CLASS_BOX_WIDTH)] 329 | + [ 330 | get_text_width( 331 | image, class_name.capitalize(), text_size, MONTSERATT_MEDIUM_TTF_PATH, 332 | ) 333 | + text_offset_x * 2 334 | for class_name in classes 335 | ], 336 | ) 337 | for class_number, class_name in enumerate(classes): 338 | color = ( 339 | picked_color if class_number == picked_class_number else BACKGROUND_COLOR 340 | ) 341 | image = add_class_box( 342 | image, color, class_name.capitalize(), box_w, class_number, 343 | ) 344 | return image 345 | 346 | 347 | def add_legend_all_classes( 348 | image: np.ndarray, classes: List, colors: List, 349 | ) -> np.ndarray: 350 | img_h, img_w, _ = image.shape 351 | scale = min([img_w, img_h]) / NORM_HEIGHT 352 | text_size = int(scale * CLASS_TEXT_SIZE) 353 | text_offset_x = int(scale * CLASS_BOX_TEXT_OFFSET_X) 354 | box_w = max( 355 | [int(scale * CLASS_BOX_WIDTH)] 356 | + [ 357 | get_text_width( 358 | image, class_name.capitalize(), text_size, MONTSERATT_MEDIUM_TTF_PATH, 359 | ) 360 | + text_offset_x * 2 361 | for class_name in classes 362 | ], 363 | ) 364 | for class_number, class_name in enumerate(classes): 365 | image = add_class_box( 366 | image, colors[class_number], class_name.capitalize(), box_w, class_number, 367 | ) 368 | return image 369 | 370 | 371 | def draw_detections_one_frame(image: np.ndarray, detections: List[BBox]) -> np.ndarray: 372 | classes = defaultdict(list) 373 | for det in detections: 374 | classes[det.class_name].append(det) 375 | for class_number, class_detections in enumerate(classes.values()): 376 | color = RGBA_COLORS[class_number] 377 | for detection in class_detections: 378 | image = add_bbox( 379 | image, [detection.x1, detection.y1, detection.x2, detection.y2], color, 380 | ) 381 | image = add_legend_all_classes(image, classes, colors=RGBA_COLORS) 382 | return image 383 | 384 | 385 | def draw_detections(image: np.ndarray, detections: List[BBox]) -> List[np.ndarray]: 386 | source_image = image.copy() 387 | images = [] 388 | classes = defaultdict(list) 389 | for det in detections: 390 | classes[det.class_name].append(det) 391 | for class_number, class_detections in enumerate(classes.values()): 392 | one_class_image = source_image.copy() 393 | color = RGBA_COLORS[class_number] 394 | for detection in class_detections: 395 | one_class_image = add_bbox( 396 | one_class_image, 397 | [detection.x1, detection.y1, detection.x2, detection.y2], 398 | color, 399 | ) 400 | one_class_image = add_legend(one_class_image, classes, color, class_number) 401 | images.append(one_class_image) 402 | images.append(add_legend(source_image, classes, BACKGROUND_COLOR, -1)) 403 | return images 404 | 405 | 406 | def draw_segmentation_one_frame(image: np.ndarray, detection: Mask) -> np.ndarray: 407 | classes = [ 408 | class_name 409 | for class_number, class_name in enumerate(detection.classes) 410 | if class_number in detection.mask["classes"] 411 | ] 412 | for class_number, rle_mask in enumerate(detection.mask["binary"]): 413 | color = RGBA_COLORS[class_number] 414 | decoded_mask = decode_coco_rle(rle_mask) 415 | idx = decoded_mask == 1 416 | image = add_mask(image, idx, color) 417 | image = add_legend_all_classes(image, classes, colors=RGBA_COLORS) 418 | return image 419 | 420 | 421 | def draw_instance_segmentation_one_frame( 422 | image: np.ndarray, instance_mask: InstanceMask, 423 | ) -> np.ndarray: 424 | if not len(instance_mask.masks): 425 | return image 426 | predicted_classes = [mask.mask["classes"] for mask in instance_mask.masks] 427 | unique_predicted_classes = list( 428 | set(elem for sublist in predicted_classes for elem in sublist), 429 | ) 430 | classes = [ 431 | class_name 432 | for class_number, class_name in enumerate(instance_mask.masks[0].classes) 433 | if class_number in unique_predicted_classes 434 | ] 435 | color_boxes = [] 436 | result_mask = np.zeros_like(image).astype(np.uint8) 437 | for instance_number, (mask, box) in enumerate( 438 | zip(instance_mask.masks, instance_mask.detections), 439 | ): 440 | # skip background here 441 | rle_mask = mask.mask["binary"][-1] 442 | color = RGBA_COLORS[instance_number] 443 | decoded_mask = decode_coco_rle(rle_mask) 444 | idx = decoded_mask == 1 445 | result_mask[idx] = color[: result_mask.shape[2]] 446 | color_boxes.append((box, color)) 447 | image = add_instance_mask(image, result_mask) 448 | for (box, color) in color_boxes: 449 | image = add_bbox(image, [box.x1, box.y1, box.x2, box.y2], color) 450 | image = add_legend_all_classes( 451 | image, classes, colors=[BACKGROUND_COLOR] * len(classes), 452 | ) 453 | return image 454 | 455 | 456 | def draw_instance_segmentation( 457 | image: np.ndarray, instance_mask: InstanceMask, 458 | ) -> np.ndarray: 459 | if not len(instance_mask.masks): 460 | return image 461 | source_image = image.copy() 462 | images = [] 463 | predicted_classes = [mask.mask["classes"] for mask in instance_mask.masks] 464 | unique_predicted_classes = list( 465 | set(elem for sublist in predicted_classes for elem in sublist), 466 | ) 467 | classes = [ 468 | class_name 469 | for class_number, class_name in enumerate(instance_mask.masks[0].classes) 470 | if class_number in unique_predicted_classes 471 | ] 472 | 473 | class_name_to_idx_mapping = dict(zip(classes, unique_predicted_classes)) 474 | 475 | per_class_masks = { 476 | class_id: np.zeros_like(image).astype(np.uint8) 477 | for class_id in unique_predicted_classes 478 | } 479 | per_class_color_boxes = {class_id: [] for class_id in unique_predicted_classes} 480 | # we should group instances by class here 481 | for instance_number, (mask, box) in enumerate( 482 | zip(instance_mask.masks, instance_mask.detections), 483 | ): 484 | rle_mask = mask.mask["binary"][-1] 485 | class_id = class_name_to_idx_mapping[box.class_name] 486 | # we loop RGBA colors if instance number grater than available colors 487 | color = RGBA_COLORS[instance_number % len(RGBA_COLORS)] 488 | decoded_mask = decode_coco_rle(rle_mask) 489 | idx = decoded_mask == 1 490 | per_class_masks[class_id][idx] = color[: per_class_masks[class_id].shape[2]] 491 | per_class_color_boxes[class_id].append((box, color)) 492 | 493 | for idx, (class_mask, color_boxes) in enumerate( 494 | zip(per_class_masks.values(), per_class_color_boxes.values()), 495 | ): 496 | one_class_image = source_image.copy() 497 | one_class_image = add_instance_mask(one_class_image, class_mask) 498 | for (box, color) in color_boxes: 499 | one_class_image = add_bbox( 500 | one_class_image, [box.x1, box.y1, box.x2, box.y2], color, 501 | ) 502 | one_class_image = add_legend(one_class_image, classes, DARK_PINK_COLOR, idx) 503 | images.append(one_class_image) 504 | images.append(add_legend(source_image, classes, BACKGROUND_COLOR, -1)) 505 | return images 506 | 507 | 508 | def draw_segmentation(image: np.ndarray, detection: Mask) -> List[np.ndarray]: 509 | source_image = image.copy() 510 | images = [] 511 | classes = [ 512 | class_name 513 | for class_number, class_name in enumerate(detection.classes) 514 | if class_number in detection.mask["classes"] 515 | ] 516 | for class_number, rle_mask in enumerate(detection.mask["binary"]): 517 | one_class_image = source_image.copy() 518 | color = RGBA_COLORS[class_number] 519 | decoded_mask = decode_coco_rle(rle_mask) 520 | idx = decoded_mask == 1 521 | one_class_image = add_mask(one_class_image, idx, color) 522 | one_class_image = add_legend(one_class_image, classes, color, class_number) 523 | images.append(one_class_image) 524 | images.append(add_legend(source_image, classes, BACKGROUND_COLOR, -1)) 525 | return images 526 | 527 | 528 | def draw_classification_one_frame(image: np.ndarray, labels: List[Label]) -> np.ndarray: 529 | labels = [label for label in labels if label.score > 0.01] 530 | img_h, img_w, _ = image.shape 531 | scale = min([img_w, img_h]) / NORM_HEIGHT 532 | box_offset_x = box_offset_y = int(scale * BOX_CORNER_OFFSET) 533 | box_h = int(scale * PROB_BOX_HEIGHT) 534 | text_offset_x = int(scale * PROB_COMMENT_OFFSET_X) 535 | text_size = int(scale * COMMENT_TEXT_SIZE) 536 | box_w = max( 537 | [int(scale * PROB_BOX_WIDTH)] 538 | + [ 539 | get_text_width( 540 | image, 541 | label.class_name.capitalize(), 542 | text_size, 543 | MONTSERATT_REGULAR_TTF_PATH, 544 | ) 545 | + text_offset_x * 2 546 | for label in labels 547 | ], 548 | ) 549 | for label_number, label in enumerate(labels): 550 | delimiter = label_number != len(labels) - 1 551 | box_x1 = img_w - box_offset_x - box_w 552 | box_y1 = box_offset_y + label_number * box_h 553 | box_x2 = box_x1 + box_w 554 | box_y2 = box_y1 + box_h 555 | coords = [box_x1, box_y1, box_x2, box_y2] 556 | image = add_probability_box( 557 | image, 558 | "{}%".format(int(label.score * 100)), 559 | label.class_name.capitalize(), 560 | coords, 561 | delimiter, 562 | ) 563 | return image 564 | 565 | 566 | def draw_classification(image: np.ndarray, labels: List[Label]) -> List[np.ndarray]: 567 | source_image = image.copy() 568 | return [draw_classification_one_frame(image, labels), source_image] 569 | 570 | 571 | def draw_text_detections_one_frame( 572 | image: np.ndarray, detections: List[TextPolygon], 573 | ) -> np.ndarray: 574 | for polygon in detections: 575 | image = add_poly(image, polygon, LAVANDER_COLOR, LINE_THINKNESS) 576 | if len(polygon.text.strip()): 577 | image = add_info( 578 | image, 579 | [polygon.points[0].x, polygon.points[0].y], 580 | LAVANDER_COLOR, 581 | polygon.text, 582 | BLACK_TEXT_COLOR, 583 | ) 584 | return image 585 | 586 | 587 | def draw_text_detections( 588 | image: np.ndarray, detections: List[TextPolygon], 589 | ) -> List[np.ndarray]: 590 | source_image = image.copy() 591 | return [draw_text_detections_one_frame(image, detections), source_image] 592 | 593 | 594 | def draw_landmarks_one_frame( 595 | image: np.ndarray, detections: List[Landmarks], draw_bbox: bool = True, 596 | ) -> np.ndarray: 597 | for detection in detections: 598 | for keypoint_number, keypoint in enumerate(detection.keypoints): 599 | image = add_point(image, keypoint, RGBA_COLORS[keypoint_number]) 600 | if draw_bbox: 601 | image = add_bbox( 602 | image, 603 | ( 604 | detection.bbox.x1, 605 | detection.bbox.y1, 606 | detection.bbox.x2, 607 | detection.bbox.y2, 608 | ), 609 | LAVANDER_COLOR, 610 | ) 611 | return image 612 | 613 | 614 | def draw_landmarks( 615 | image: np.ndarray, detections: List[Landmarks], draw_bbox: bool = True, 616 | ) -> List[np.ndarray]: 617 | source_image = image.copy() 618 | return [draw_landmarks_one_frame(image, detections, draw_bbox), source_image] 619 | 620 | 621 | def draw_keypoints_one_frame(image: np.ndarray, detections: List[Pose]) -> np.ndarray: 622 | for detection in detections: 623 | image = add_pose(image, detection, LAVANDER_COLOR) 624 | return image 625 | 626 | 627 | def draw_keypoints(image: np.ndarray, detections: List[Pose]) -> List[np.ndarray]: 628 | source_image = image.copy() 629 | images = [] 630 | for detection in detections: 631 | images.append(draw_keypoints_one_frame(image.copy(), [detection])) 632 | images.append(draw_keypoints_one_frame(image, detections)) 633 | images.append(source_image) 634 | return images 635 | 636 | 637 | def draw_mesh_one_frame( 638 | image: np.ndarray, detections: List[Landmarks], draw_bbox: bool = True, 639 | ) -> np.ndarray: 640 | for detection in detections: 641 | for keypoint in detection.keypoints: 642 | image = add_point(image, keypoint, DARK_PINK_COLOR, MESH_POINT_SIZE) 643 | if draw_bbox: 644 | image = add_bbox( 645 | image, 646 | ( 647 | detection.bbox.x1, 648 | detection.bbox.y1, 649 | detection.bbox.x2, 650 | detection.bbox.y2, 651 | ), 652 | LAVANDER_COLOR, 653 | ) 654 | return image 655 | 656 | 657 | def draw_mesh( 658 | image: np.ndarray, detections: List[Landmarks], draw_bbox: bool = True, 659 | ) -> List[np.ndarray]: 660 | source_image = image.copy() 661 | return [draw_mesh_one_frame(image, detections, draw_bbox), source_image] 662 | 663 | 664 | def draw_age_gender_recognition_one_frame( 665 | image: np.ndarray, detections: List[AgeGenderLabel], 666 | ) -> np.ndarray: 667 | for detection in detections: 668 | coords = max(detection.bbox.x1, 0), max(detection.bbox.y1, 0) 669 | gender = max(detection.genders, key=lambda x: x.score).class_name 670 | text = "{}, {} y.o.".format( 671 | "MEN" if gender == "male" else "WOMAN", detection.age, 672 | ) 673 | image = add_info(image, coords, BACKGROUND_COLOR, text, WHITE_TEXT_COLOR) 674 | return image 675 | 676 | 677 | def draw_age_gender_recognition( 678 | image: np.ndarray, detections: List[AgeGenderLabel], 679 | ) -> List[np.ndarray]: 680 | source_image = image.copy() 681 | images = [] 682 | for detection in detections: 683 | images.append(draw_age_gender_recognition_one_frame(image.copy(), [detection])) 684 | images.append(draw_age_gender_recognition_one_frame(image.copy(), detections)) 685 | images.append(source_image) 686 | return images 687 | 688 | 689 | def draw_emotion_recognition_one_frame( 690 | image: np.ndarray, detections: List[EmotionLabel], 691 | ) -> np.ndarray: 692 | scale = min(image.shape[:2]) / NORM_HEIGHT 693 | box_h = int(scale * PROB_BOX_HEIGHT) 694 | for label in detections: 695 | emotion = max(label.emotions, key=lambda x: x.score) 696 | x = max(label.bbox.x1, 0) 697 | y = max(label.bbox.y1 - box_h, 0) 698 | box_w = max( 699 | int(scale * PROB_BOX_WIDTH), 700 | get_text_width( 701 | image, 702 | emotion.class_name.capitalize(), 703 | int(scale * COMMENT_TEXT_SIZE), 704 | MONTSERATT_REGULAR_TTF_PATH, 705 | ) 706 | + int(scale * PROB_COMMENT_OFFSET_X) * 2, 707 | ) 708 | image = add_probability_box( 709 | image, 710 | "{}%".format(int(emotion.score * 100)), 711 | emotion.class_name.capitalize(), 712 | [x, y, x + box_w, y + box_h], 713 | False, 714 | ) 715 | return image 716 | 717 | 718 | def draw_emotion_recognition( 719 | image: np.ndarray, detections: List[EmotionLabel], 720 | ) -> List[np.ndarray]: 721 | images = [] 722 | source_image = image.copy() 723 | for detection in detections: 724 | images.append(draw_emotion_recognition_one_frame(image.copy(), [detection])) 725 | images.append(draw_emotion_recognition_one_frame(image, detections)) 726 | images.append(source_image) 727 | return images 728 | 729 | 730 | def draw_tracks_one_frame(image: np.ndarray, detection: VideoFrame) -> np.ndarray: 731 | for box in detection.boxes: 732 | image = add_bbox(image, [box.x1, box.y1, box.x2, box.y2], LAVANDER_COLOR) 733 | image = add_info( 734 | image, 735 | (box.x1, box.y1), 736 | LAVANDER_COLOR, 737 | f"id:{box.track_number}", 738 | BLACK_TEXT_COLOR, 739 | ) 740 | return image 741 | 742 | 743 | def draw_tracks( 744 | video: Generator, frames: List[VideoFrame], save_path: str, fps: int, 745 | ) -> None: 746 | outputdict = FFMPEG_OUTPUT_DICT.copy() 747 | outputdict["-r"] = str(fps) 748 | writer = skvideo.io.FFmpegWriter(save_path, outputdict=outputdict) 749 | for frame_number, frame in enumerate(video): 750 | if frame_number >= len(frames): 751 | continue 752 | anno = frames[frame_number] 753 | frame = draw_tracks_one_frame(frame, anno) 754 | writer.writeFrame(frame) 755 | writer.close() 756 | 757 | 758 | def draw_countable_tracks_one_frame( 759 | image: np.ndarray, detection: CountableVideoFrame, 760 | ) -> np.ndarray: 761 | for box in detection.boxes: 762 | image = add_bbox(image, [box.x1, box.y1, box.x2, box.y2], LAVANDER_COLOR) 763 | image = add_info( 764 | image, 765 | (box.x1, box.y1), 766 | LAVANDER_COLOR, 767 | f"id:{box.track_number}", 768 | BLACK_TEXT_COLOR, 769 | ) 770 | img_h, img_w = image.shape[:2] 771 | scale = min(img_w, img_h) / NORM_HEIGHT 772 | box_offset_x = box_offset_y = int(scale * BOX_CORNER_OFFSET) 773 | box_h = box_w = int(scale * PROB_BOX_WIDTH) 774 | if detection.people_in is not None: 775 | image = add_probability_box( 776 | image, 777 | str(detection.people_in), 778 | "Came", 779 | [ 780 | img_w - box_offset_x - box_w * 2, 781 | box_offset_y, 782 | img_w - box_offset_x - box_w, 783 | box_offset_y + box_h, 784 | ], 785 | False, 786 | CORALL_COLOR, 787 | ) 788 | if detection.people_out is not None: 789 | image = add_probability_box( 790 | image, 791 | str(detection.people_out), 792 | "Left", 793 | [ 794 | img_w - box_offset_x - box_w, 795 | box_offset_y, 796 | img_w - box_offset_x, 797 | box_offset_y + box_h, 798 | ], 799 | False, 800 | BACKGROUND_COLOR, 801 | ) 802 | return image 803 | 804 | 805 | def draw_countable_tracks( 806 | video: Generator, frames: List[CountableVideoFrame], save_path: str, fps: int, 807 | ) -> None: 808 | outputdict = FFMPEG_OUTPUT_DICT.copy() 809 | outputdict["-r"] = str(fps) 810 | writer = skvideo.io.FFmpegWriter(save_path, outputdict=outputdict) 811 | for frame_number, frame in enumerate(video): 812 | if frame_number >= len(frames): 813 | continue 814 | anno = frames[frame_number] 815 | frame = draw_countable_tracks_one_frame(frame, anno) 816 | writer.writeFrame(frame) 817 | writer.close() 818 | 819 | 820 | def create_gif(images: List[np.ndarray], save_path: str, fps: int = 1) -> None: 821 | imageio.mimsave(save_path, images, format="GIF-FI", fps=fps, quantizer="nq") 822 | 823 | 824 | def create_image(image: np.ndarray, save_path: str) -> None: 825 | cv2.imwrite(save_path, image) 826 | 827 | 828 | def draw_background_removal(image: np.ndarray, detection: Mask) -> np.ndarray: 829 | foreground_mask_id = detection.classes.index("foreground") 830 | if len(detection.mask["binary"]) > foreground_mask_id: 831 | decoded_mask = decode_coco_rle(detection.mask["binary"][foreground_mask_id]) 832 | idx = decoded_mask == 0 833 | image[idx] = (255, 255, 255) 834 | alfa_channel = decoded_mask.astype(image.dtype)[:, :, np.newaxis] * 255 835 | else: 836 | alfa_channel = np.ones((image.shape[0], image.shape[1], 1), dtype=image.dtype) 837 | 838 | image = np.concatenate((image, alfa_channel), axis=2) 839 | return image 840 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.black] 2 | line-length = 88 3 | target-version = ['py36', 'py37', 'py38'] 4 | include = '\.pyi?$' 5 | exclude = ''' 6 | /( 7 | \.eggs 8 | | \.git 9 | | \.hg 10 | | \.mypy_cache 11 | | \.tox 12 | | \.venv 13 | | _build 14 | | buck-out 15 | | build 16 | | dist 17 | )/ 18 | ''' 19 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | black == 19.10b0 2 | flake8 == 3.8.2 3 | flake8-bugbear == 20.1.4 4 | isort[pyproject] == 4.3.21 5 | pre-commit == 2.5.1 6 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import codecs 4 | import os.path 5 | 6 | from setuptools import setup 7 | 8 | 9 | def get_version(rel_path): 10 | with codecs.open( 11 | os.path.join(os.path.abspath(os.path.dirname(__file__)), rel_path), "r", 12 | ) as fp: 13 | for line in fp.read().splitlines(): 14 | if line.startswith("__version__"): 15 | delim = '"' if '"' in line else "'" 16 | return line.split(delim)[1] 17 | raise RuntimeError("Unable to find version string.") 18 | 19 | 20 | packages = ["modelplace_api"] 21 | 22 | package_data = {"": ["*", "text_styles/*"]} 23 | install_requires = [ 24 | "pydantic>=1.5.1,<1.9.0", 25 | "loguru>=0.5.1", 26 | "numpy>=1.16.4", 27 | ] 28 | 29 | extras_require = { 30 | "vis": [ 31 | "Pillow>=7.1.2", 32 | "opencv-python>=4.2.0.34,<5.0", 33 | "imageio==2.9.0", 34 | "sk-video==1.1.10", 35 | "pycocotools==2.0.2", 36 | ], 37 | "vis-windows": [ 38 | "Pillow>=7.1.2", 39 | "opencv-python>=4.2.0.34,<5.0", 40 | "imageio==2.9.0", 41 | "sk-video==1.1.10", 42 | ], 43 | } 44 | 45 | setup_kwargs = { 46 | "name": "modelplace-api", 47 | "version": get_version("modelplace_api/__init__.py"), 48 | "description": "", 49 | "long_description": None, 50 | "author": "OpenCV.AI", 51 | "author_email": "modelplace@opencv.ai", 52 | "maintainer": "OpenCV.AI", 53 | "maintainer_email": "modelplace@opencv.ai", 54 | "url": None, 55 | "packages": packages, 56 | "package_data": package_data, 57 | "install_requires": install_requires, 58 | "python_requires": ">=3.6,<4.0", 59 | "extras_require": extras_require, 60 | } 61 | 62 | setup(**setup_kwargs) 63 | -------------------------------------------------------------------------------- /tests/decoded_binary_mask.json: -------------------------------------------------------------------------------- 1 | {"size": [4000, 3000], "counts": "lfSj31ck3h2QNV1oNl0VOh0ZOd0\\Od0]Ob0_O`0A>A`0BC@?B?@a0^Oc0\\Od0\\Of0XOj0TOQ1jNScN"} -------------------------------------------------------------------------------- /tests/test_rle.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | 4 | import cv2 5 | import numpy as np 6 | 7 | from modelplace_api._rle_mask import decode_coco_rle, encode_binary_mask 8 | 9 | binary_mask = np.zeros((4000, 3000)) 10 | binary_mask = np.asfortranarray( 11 | cv2.circle(binary_mask, (2000, 1500), 1000, (255, 255, 255), -1) // 255, 12 | ).astype(np.uint8) 13 | 14 | 15 | with open("decoded_binary_mask.json", "r") as f: 16 | encoded_binary_mask = json.load(f) 17 | encoded_binary_mask["counts"] = encoded_binary_mask["counts"].encode("utf-8") 18 | 19 | 20 | def test_encode(): 21 | python_encoded_mask = encode_binary_mask(binary_mask) 22 | assert python_encoded_mask == encoded_binary_mask 23 | 24 | 25 | def test_decode(): 26 | python_decoded_mask = decode_coco_rle(encoded_binary_mask) 27 | assert np.all(python_decoded_mask == binary_mask) 28 | -------------------------------------------------------------------------------- /tests/test_rle_backward_compatibility.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import cv2 4 | import numpy as np 5 | from pycocotools import mask 6 | 7 | from modelplace_api._rle_mask import decode_coco_rle, encode_binary_mask 8 | 9 | binary_mask = np.zeros((4000, 3000)) 10 | binary_mask = np.asfortranarray( 11 | cv2.circle(binary_mask, (2000, 1500), 1000, (255, 255, 255), -1) // 255, 12 | ).astype(np.uint8) 13 | 14 | 15 | def test_encode(): 16 | pycocotools_encoded_mask = mask.encode(binary_mask) 17 | python_encoded_mask = encode_binary_mask(binary_mask) 18 | assert python_encoded_mask == pycocotools_encoded_mask 19 | 20 | 21 | def test_decode(): 22 | pycocotools_encoded_mask = mask.encode(binary_mask) 23 | python_decoded_mask = decode_coco_rle(pycocotools_encoded_mask) 24 | assert np.all(python_decoded_mask == binary_mask) 25 | --------------------------------------------------------------------------------