Init

.gitignore

@@ -0,0 +1,162 @@
+# 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/

anime_face_detector/__init__.py

@@ -0,0 +1,54 @@
+import pathlib
+
+import torch
+
+from .detector import LandmarkDetector
+
+
+def get_config_path(model_name: str) -> pathlib.Path:
+    assert model_name in ['faster-rcnn', 'yolov3', 'hrnetv2']
+
+    package_path = pathlib.Path(__file__).parent.resolve()
+    if model_name in ['faster-rcnn', 'yolov3']:
+        config_dir = package_path / 'configs' / 'mmdet'
+    else:
+        config_dir = package_path / 'configs' / 'mmpose'
+    return config_dir / f'{model_name}.py'
+
+
+def get_checkpoint_path(model_name: str) -> pathlib.Path:
+    assert model_name in ['faster-rcnn', 'yolov3', 'hrnetv2']
+    if model_name in ['faster-rcnn', 'yolov3']:
+        file_name = f'mmdet_anime-face_{model_name}.pth'
+    else:
+        file_name = f'mmpose_anime-face_{model_name}.pth'
+
+    model_dir = pathlib.Path(torch.hub.get_dir()) / 'checkpoints'
+    model_dir.mkdir(exist_ok=True, parents=True)
+    model_path = model_dir / file_name
+    if not model_path.exists():
+        url = f'https://github.com/hysts/anime-face-detector/releases/download/v0.0.1/{file_name}'
+        torch.hub.download_url_to_file(url, model_path.as_posix())
+
+    return model_path
+
+
+def create_detector(face_detector_name: str = 'yolov3',
+                    landmark_model_name='hrnetv2',
+                    device: str = 'cuda:0',
+                    flip_test: bool = True,
+                    box_scale_factor: float = 1.1) -> LandmarkDetector:
+    assert face_detector_name in ['yolov3', 'faster-rcnn']
+    assert landmark_model_name in ['hrnetv2']
+    detector_config_path = get_config_path(face_detector_name)
+    landmark_config_path = get_config_path(landmark_model_name)
+    detector_checkpoint_path = get_checkpoint_path(face_detector_name)
+    landmark_checkpoint_path = get_checkpoint_path(landmark_model_name)
+    model = LandmarkDetector(landmark_config_path,
+                             landmark_checkpoint_path,
+                             detector_config_path,
+                             detector_checkpoint_path,
+                             device=device,
+                             flip_test=flip_test,
+                             box_scale_factor=box_scale_factor)
+    return model

anime_face_detector/configs/mmdet/faster-rcnn.py

@@ -0,0 +1,66 @@
+model = dict(type='FasterRCNN',
+             backbone=dict(type='ResNet',
+                           depth=50,
+                           num_stages=4,
+                           out_indices=(0, 1, 2, 3),
+                           frozen_stages=1,
+                           norm_cfg=dict(type='BN', requires_grad=True),
+                           norm_eval=True,
+                           style='pytorch'),
+             neck=dict(type='FPN',
+                       in_channels=[256, 512, 1024, 2048],
+                       out_channels=256,
+                       num_outs=5),
+             rpn_head=dict(type='RPNHead',
+                           in_channels=256,
+                           feat_channels=256,
+                           anchor_generator=dict(type='AnchorGenerator',
+                                                 scales=[8],
+                                                 ratios=[0.5, 1.0, 2.0],
+                                                 strides=[4, 8, 16, 32, 64]),
+                           bbox_coder=dict(type='DeltaXYWHBBoxCoder',
+                                           target_means=[0.0, 0.0, 0.0, 0.0],
+                                           target_stds=[1.0, 1.0, 1.0, 1.0])),
+             roi_head=dict(
+                 type='StandardRoIHead',
+                 bbox_roi_extractor=dict(type='SingleRoIExtractor',
+                                         roi_layer=dict(type='RoIAlign',
+                                                        output_size=7,
+                                                        sampling_ratio=0),
+                                         out_channels=256,
+                                         featmap_strides=[4, 8, 16, 32]),
+                 bbox_head=dict(type='Shared2FCBBoxHead',
+                                in_channels=256,
+                                fc_out_channels=1024,
+                                roi_feat_size=7,
+                                num_classes=1,
+                                bbox_coder=dict(
+                                    type='DeltaXYWHBBoxCoder',
+                                    target_means=[0.0, 0.0, 0.0, 0.0],
+                                    target_stds=[0.1, 0.1, 0.2, 0.2]),
+                                reg_class_agnostic=False)),
+             test_cfg=dict(rpn=dict(nms_pre=1000,
+                                    max_per_img=1000,
+                                    nms=dict(type='nms', iou_threshold=0.7),
+                                    min_bbox_size=0),
+                           rcnn=dict(score_thr=0.05,
+                                     nms=dict(type='nms', iou_threshold=0.5),
+                                     max_per_img=100)))
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='MultiScaleFlipAug',
+         img_scale=(1333, 800),
+         flip=False,
+         transforms=[
+             dict(type='Resize', keep_ratio=True),
+             dict(type='RandomFlip'),
+             dict(type='Normalize',
+                  mean=[123.675, 116.28, 103.53],
+                  std=[58.395, 57.12, 57.375],
+                  to_rgb=True),
+             dict(type='Pad', size_divisor=32),
+             dict(type='DefaultFormatBundle'),
+             dict(type='Collect', keys=['img'])
+         ])
+]
+data = dict(test=dict(pipeline=test_pipeline))

anime_face_detector/configs/mmdet/yolov3.py

@@ -0,0 +1,47 @@
+model = dict(type='YOLOV3',
+             backbone=dict(type='Darknet', depth=53, out_indices=(3, 4, 5)),
+             neck=dict(type='YOLOV3Neck',
+                       num_scales=3,
+                       in_channels=[1024, 512, 256],
+                       out_channels=[512, 256, 128]),
+             bbox_head=dict(type='YOLOV3Head',
+                            num_classes=1,
+                            in_channels=[512, 256, 128],
+                            out_channels=[1024, 512, 256],
+                            anchor_generator=dict(type='YOLOAnchorGenerator',
+                                                  base_sizes=[[(116, 90),
+                                                               (156, 198),
+                                                               (373, 326)],
+                                                              [(30, 61),
+                                                               (62, 45),
+                                                               (59, 119)],
+                                                              [(10, 13),
+                                                               (16, 30),
+                                                               (33, 23)]],
+                                                  strides=[32, 16, 8]),
+                            bbox_coder=dict(type='YOLOBBoxCoder'),
+                            featmap_strides=[32, 16, 8]),
+             test_cfg=dict(nms_pre=1000,
+                           min_bbox_size=0,
+                           score_thr=0.05,
+                           conf_thr=0.005,
+                           nms=dict(type='nms', iou_threshold=0.45),
+                           max_per_img=100))
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='MultiScaleFlipAug',
+         img_scale=(608, 608),
+         flip=False,
+         transforms=[
+             dict(type='Resize', keep_ratio=True),
+             dict(type='RandomFlip'),
+             dict(type='Normalize',
+                  mean=[0, 0, 0],
+                  std=[255.0, 255.0, 255.0],
+                  to_rgb=True),
+             dict(type='Pad', size_divisor=32),
+             dict(type='DefaultFormatBundle'),
+             dict(type='Collect', keys=['img'])
+         ])
+]
+data = dict(test=dict(pipeline=test_pipeline))

anime_face_detector/configs/mmpose/hrnetv2.py

@@ -0,0 +1,250 @@
+channel_cfg = dict(num_output_channels=28,
+                   dataset_joints=28,
+                   dataset_channel=[
+                       list(range(28)),
+                   ],
+                   inference_channel=list(range(28)))
+
+model = dict(
+    type='TopDown',
+    backbone=dict(type='HRNet',
+                  in_channels=3,
+                  extra=dict(stage1=dict(num_modules=1,
+                                         num_branches=1,
+                                         block='BOTTLENECK',
+                                         num_blocks=(4, ),
+                                         num_channels=(64, )),
+                             stage2=dict(num_modules=1,
+                                         num_branches=2,
+                                         block='BASIC',
+                                         num_blocks=(4, 4),
+                                         num_channels=(18, 36)),
+                             stage3=dict(num_modules=4,
+                                         num_branches=3,
+                                         block='BASIC',
+                                         num_blocks=(4, 4, 4),
+                                         num_channels=(18, 36, 72)),
+                             stage4=dict(num_modules=3,
+                                         num_branches=4,
+                                         block='BASIC',
+                                         num_blocks=(4, 4, 4, 4),
+                                         num_channels=(18, 36, 72, 144),
+                                         multiscale_output=True),
+                             upsample=dict(mode='bilinear',
+                                           align_corners=False))),
+    keypoint_head=dict(type='TopdownHeatmapSimpleHead',
+                       in_channels=[18, 36, 72, 144],
+                       in_index=(0, 1, 2, 3),
+                       input_transform='resize_concat',
+                       out_channels=channel_cfg['num_output_channels'],
+                       num_deconv_layers=0,
+                       extra=dict(final_conv_kernel=1,
+                                  num_conv_layers=1,
+                                  num_conv_kernels=(1, )),
+                       loss_keypoint=dict(type='JointsMSELoss',
+                                          use_target_weight=True)),
+    test_cfg=dict(flip_test=True,
+                  post_process='unbiased',
+                  shift_heatmap=True,
+                  modulate_kernel=11))
+
+data_cfg = dict(image_size=[256, 256],
+                heatmap_size=[64, 64],
+                num_output_channels=channel_cfg['num_output_channels'],
+                num_joints=channel_cfg['dataset_joints'],
+                dataset_channel=channel_cfg['dataset_channel'],
+                inference_channel=channel_cfg['inference_channel'])
+
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='TopDownAffine'),
+    dict(type='ToTensor'),
+    dict(type='NormalizeTensor',
+         mean=[0.485, 0.456, 0.406],
+         std=[0.229, 0.224, 0.225]),
+    dict(type='Collect',
+         keys=['img'],
+         meta_keys=['image_file', 'center', 'scale', 'rotation',
+                    'flip_pairs']),
+]
+
+dataset_info = dict(dataset_name='anime_face',
+                    paper_info=dict(),
+                    keypoint_info={
+                        0:
+                        dict(name='kpt-0',
+                             id=0,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-4'),
+                        1:
+                        dict(name='kpt-1',
+                             id=1,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-3'),
+                        2:
+                        dict(name='kpt-2',
+                             id=2,
+                             color=[255, 255, 255],
+                             type='',
+                             swap=''),
+                        3:
+                        dict(name='kpt-3',
+                             id=3,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-1'),
+                        4:
+                        dict(name='kpt-4',
+                             id=4,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-0'),
+                        5:
+                        dict(name='kpt-5',
+                             id=5,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-10'),
+                        6:
+                        dict(name='kpt-6',
+                             id=6,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-9'),
+                        7:
+                        dict(name='kpt-7',
+                             id=7,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-8'),
+                        8:
+                        dict(name='kpt-8',
+                             id=8,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-7'),
+                        9:
+                        dict(name='kpt-9',
+                             id=9,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-6'),
+                        10:
+                        dict(name='kpt-10',
+                             id=10,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-5'),
+                        11:
+                        dict(name='kpt-11',
+                             id=11,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-19'),
+                        12:
+                        dict(name='kpt-12',
+                             id=12,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-18'),
+                        13:
+                        dict(name='kpt-13',
+                             id=13,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-17'),
+                        14:
+                        dict(name='kpt-14',
+                             id=14,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-22'),
+                        15:
+                        dict(name='kpt-15',
+                             id=15,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-21'),
+                        16:
+                        dict(name='kpt-16',
+                             id=16,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-20'),
+                        17:
+                        dict(name='kpt-17',
+                             id=17,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-13'),
+                        18:
+                        dict(name='kpt-18',
+                             id=18,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-12'),
+                        19:
+                        dict(name='kpt-19',
+                             id=19,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-11'),
+                        20:
+                        dict(name='kpt-20',
+                             id=20,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-16'),
+                        21:
+                        dict(name='kpt-21',
+                             id=21,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-15'),
+                        22:
+                        dict(name='kpt-22',
+                             id=22,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-14'),
+                        23:
+                        dict(name='kpt-23',
+                             id=23,
+                             color=[255, 255, 255],
+                             type='',
+                             swap=''),
+                        24:
+                        dict(name='kpt-24',
+                             id=24,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-26'),
+                        25:
+                        dict(name='kpt-25',
+                             id=25,
+                             color=[255, 255, 255],
+                             type='',
+                             swap=''),
+                        26:
+                        dict(name='kpt-26',
+                             id=26,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='kpt-24'),
+                        27:
+                        dict(name='kpt-27',
+                             id=27,
+                             color=[255, 255, 255],
+                             type='',
+                             swap='')
+                    },
+                    skeleton_info={},
+                    joint_weights=[1.] * 28,
+                    sigmas=[])
+
+data = dict(test=dict(type='',
+                      data_cfg=data_cfg,
+                      pipeline=test_pipeline,
+                      dataset_info=dataset_info), )

anime_face_detector/detector.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import pathlib
+import warnings
+from typing import Optional, Union
+
+import cv2
+import mmcv
+import numpy as np
+import torch.nn as nn
+from mmdet.apis import inference_detector, init_detector
+from mmpose.apis import inference_top_down_pose_model, init_pose_model
+from mmpose.datasets import DatasetInfo
+
+
+class LandmarkDetector:
+    def __init__(
+            self,
+            landmark_detector_config_or_path: Union[mmcv.Config, str,
+                                                    pathlib.Path],
+            landmark_detector_checkpoint_path: Union[str, pathlib.Path],
+            face_detector_config_or_path: Optional[Union[mmcv.Config, str,
+                                                         pathlib.Path]] = None,
+            face_detector_checkpoint_path: Optional[Union[
+                str, pathlib.Path]] = None,
+            device: str = 'cuda:0',
+            flip_test: bool = True,
+            box_scale_factor: float = 1.1):
+        landmark_config = self._load_config(landmark_detector_config_or_path)
+        self.dataset_info = DatasetInfo(
+            landmark_config.dataset_info)  # type: ignore
+        face_detector_config = self._load_config(face_detector_config_or_path)
+
+        self.landmark_detector = self._init_pose_model(
+            landmark_config, landmark_detector_checkpoint_path, device,
+            flip_test)
+        self.face_detector = self._init_face_detector(
+            face_detector_config, face_detector_checkpoint_path, device)
+
+        self.box_scale_factor = box_scale_factor
+
+    @staticmethod
+    def _load_config(
+        config_or_path: Optional[Union[mmcv.Config, str, pathlib.Path]]
+    ) -> Optional[mmcv.Config]:
+        if config_or_path is None or isinstance(config_or_path, mmcv.Config):
+            return config_or_path
+        return mmcv.Config.fromfile(config_or_path)
+
+    @staticmethod
+    def _init_pose_model(config: mmcv.Config,
+                         checkpoint_path: Union[str, pathlib.Path],
+                         device: str, flip_test: bool) -> nn.Module:
+        if isinstance(checkpoint_path, pathlib.Path):
+            checkpoint_path = checkpoint_path.as_posix()
+        model = init_pose_model(config, checkpoint_path, device=device)
+        model.cfg.model.test_cfg.flip_test = flip_test
+        return model
+
+    @staticmethod
+    def _init_face_detector(config: Optional[mmcv.Config],
+                            checkpoint_path: Optional[Union[str,
+                                                            pathlib.Path]],
+                            device: str) -> Optional[nn.Module]:
+        if config is not None:
+            if isinstance(checkpoint_path, pathlib.Path):
+                checkpoint_path = checkpoint_path.as_posix()
+            model = init_detector(config, checkpoint_path, device=device)
+        else:
+            model = None
+        return model
+
+    def _detect_faces(self, image: np.ndarray) -> list[np.ndarray]:
+        # predicted boxes using mmdet model have the format of
+        # [x0, y0, x1, y1, score]
+        boxes = inference_detector(self.face_detector, image)[0]
+        # scale boxes by `self.box_scale_factor`
+        boxes = self._update_pred_box(boxes)
+        return boxes
+
+    def _update_pred_box(self, pred_boxes: np.ndarray) -> list[np.ndarray]:
+        boxes = []
+        for pred_box in pred_boxes:
+            box = pred_box[:4]
+            size = box[2:] - box[:2] + 1
+            new_size = size * self.box_scale_factor
+            center = (box[:2] + box[2:]) / 2
+            tl = center - new_size / 2
+            br = tl + new_size
+            pred_box[:4] = np.concatenate([tl, br])
+            boxes.append(pred_box)
+        return boxes
+
+    def _detect_landmarks(
+            self, image: np.ndarray,
+            boxes: list[dict[str, np.ndarray]]) -> list[dict[str, np.ndarray]]:
+        preds, _ = inference_top_down_pose_model(
+            self.landmark_detector,
+            image,
+            boxes,
+            format='xyxy',
+            dataset_info=self.dataset_info,
+            return_heatmap=False)
+        return preds
+
+    @staticmethod
+    def _load_image(
+            image_or_path: Union[np.ndarray, str, pathlib.Path]) -> np.ndarray:
+        if isinstance(image_or_path, np.ndarray):
+            image = image_or_path
+        elif isinstance(image_or_path, str):
+            image = cv2.imread(image_or_path)
+        elif isinstance(image_or_path, pathlib.Path):
+            image = cv2.imread(image_or_path.as_posix())
+        else:
+            raise ValueError
+        return image
+
+    def __call__(
+        self,
+        image_or_path: Union[np.ndarray, str, pathlib.Path],
+        boxes: Optional[list[np.ndarray]] = None
+    ) -> list[dict[str, np.ndarray]]:
+        """Detect face landmarks.
+
+        Args:
+            image_or_path: An image with BGR channel order or an image path.
+            boxes: A list of bounding boxes for faces. Each bounding box
+                should be of the form [x0, y0, x1, y1, [score]].
+
+        Returns: A list of detection results. Each detection result has
+            bounding box of the form [x0, y0, x1, y1, [score]], and landmarks
+            of the form [x, y, score].
+        """
+        image = self._load_image(image_or_path)
+        if boxes is None:
+            if self.face_detector is not None:
+                boxes = self._detect_faces(image)
+            else:
+                warnings.warn(
+                    'Neither the face detector nor the bounding box is '
+                    'specified. So the entire image is treated as the face '
+                    'region.')
+                h, w = image.shape[:2]
+                boxes = [np.array([0, 0, w - 1, h - 1, 1])]
+        box_list = [{'bbox': box} for box in boxes]
+        return self._detect_landmarks(image, box_list)

app.py

@@ -0,0 +1,120 @@
+import argparse
+import functools
+import pathlib
+
+import cv2
+import gradio as gr
+import numpy as np
+import PIL.Image
+import torch
+
+import anime_face_detector
+
+
+def detect(
+    img,
+    face_score_threshold: float,
+    landmark_score_threshold: float,
+    detector: anime_face_detector.LandmarkDetector,
+) -> PIL.Image.Image:
+    if not img:
+        return None
+
+    image = cv2.imread(img)
+    preds = detector(image)
+
+    res = image.copy()
+    for pred in preds:
+        box = pred["bbox"]
+        box, score = box[:4], box[4]
+        if score < face_score_threshold:
+            continue
+        box = np.round(box).astype(int)
+
+        lt = max(2, int(3 * (box[2:] - box[:2]).max() / 256))
+
+        cv2.rectangle(res, tuple(box[:2]), tuple(box[2:]), (0, 255, 0), lt)
+
+        pred_pts = pred["keypoints"]
+        for *pt, score in pred_pts:
+            if score < landmark_score_threshold:
+                color = (0, 255, 255)
+            else:
+                color = (0, 0, 255)
+            pt = np.round(pt).astype(int)
+            cv2.circle(res, tuple(pt), lt, color, cv2.FILLED)
+    res = cv2.cvtColor(res, cv2.COLOR_BGR2RGB)
+
+    image_pil = PIL.Image.fromarray(res)
+    return image_pil
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--detector", type=str, default="yolov3", choices=["yolov3", "faster-rcnn"]
+    )
+    parser.add_argument(
+        "--device", type=str, default="cuda:0", choices=["cuda:0", "cpu"]
+    )
+    parser.add_argument("--face-score-threshold", type=float, default=0.5)
+    parser.add_argument("--landmark-score-threshold", type=float, default=0.3)
+    parser.add_argument("--score-slider-step", type=float, default=0.05)
+    parser.add_argument("--port", type=int)
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--share", action="store_true")
+    parser.add_argument("--live", action="store_true")
+    args = parser.parse_args()
+
+    sample_path = pathlib.Path("assets/input.jpg")
+    if not sample_path.exists():
+        torch.hub.download_url_to_file(
+            "https://raw.githubusercontent.com/edisonlee55/hysts-anime-face-detector/main/assets/input.jpg",
+            sample_path.as_posix(),
+        )
+
+    detector = anime_face_detector.create_detector(args.detector, device=args.device)
+    func = functools.partial(detect, detector=detector)
+
+    title = "edisonlee55/hysts-anime-face-detector"
+    description = "Demo for edisonlee55/hysts-anime-face-detector. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
+    article = "<a href='https://github.com/edisonlee55/hysts-anime-face-detector'>GitHub Repo</a>"
+
+    gr.Interface(
+        func,
+        [
+            gr.Image(type="filepath", label="Input"),
+            gr.Slider(
+                0,
+                1,
+                step=args.score_slider_step,
+                value=args.face_score_threshold,
+                label="Face Score Threshold",
+            ),
+            gr.Slider(
+                0,
+                1,
+                step=args.score_slider_step,
+                value=args.landmark_score_threshold,
+                label="Landmark Score Threshold",
+            ),
+        ],
+        gr.Image(type="pil", label="Output"),
+        title=title,
+        description=description,
+        article=article,
+        examples=[
+            [
+                sample_path.as_posix(),
+                args.face_score_threshold,
+                args.landmark_score_threshold,
+            ],
+        ],
+        live=args.live,
+    ).launch(
+        debug=args.debug, share=args.share, server_port=args.port, enable_queue=True
+    )
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,16 @@
+# python=3.10
+openmim==0.3.7
+mmcv-full==1.6.2
+mmdet==2.28.2
+mmpose==0.29.0
+
+numpy==1.24.3
+scipy==1.10.1
+
+opencv-python-headless==4.7.0.72
+
+torch==2.0.1
+torchvision==0.15.2
+
+# for gradio
+# gradio==3.32.0