Spaces:
Runtime error
Runtime error
| import argparse | |
| import functools | |
| import pathlib | |
| import os | |
| import subprocess | |
| import tarfile | |
| if os.environ.get("SYSTEM") == "spaces": | |
| import mim | |
| mim.uninstall("mmcv-full", confirm_yes=True) | |
| subprocess.call("mim install mmcv-full==1.6.2".split()) | |
| subprocess.call("pip uninstall -y opencv-python".split()) | |
| subprocess.call("pip uninstall -y opencv-python-headless".split()) | |
| subprocess.call("pip install opencv-python-headless==4.7.0.72".split()) | |
| import cv2 | |
| import gradio as gr | |
| import huggingface_hub | |
| import numpy as np | |
| import PIL.Image | |
| import anime_face_detector | |
| def load_sample_image_paths(): | |
| image_dir = pathlib.Path("images") | |
| if not image_dir.exists(): | |
| dataset_repo = "hysts/sample-images-TADNE" | |
| path = huggingface_hub.hf_hub_download( | |
| dataset_repo, "images.tar.gz", repo_type="dataset" | |
| ) | |
| with tarfile.open(path) as f: | |
| f.extractall() | |
| return sorted(image_dir.glob("*")) | |
| 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="cpu", 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() | |
| image_paths = load_sample_image_paths() | |
| examples = [[path.as_posix(), 0.5, 0.3] for path in image_paths] | |
| 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=examples, | |
| live=args.live, | |
| ).launch(debug=args.debug, share=args.share, server_port=args.port) | |
| if __name__ == "__main__": | |
| main() | |