ayousanz's picture
Fix: auto-detect device for GPU/CPU fallback
88317fd
"""Anime Face Detector - Hugging Face Space with Zero GPU support"""
import subprocess
import sys
# Install mmpose without dependencies (to avoid chumpy build failure)
# Suppress output to avoid interfering with Gradio startup
subprocess.run(
[sys.executable, '-m', 'pip', 'install', '--no-deps', 'mmpose>=1.0.0'],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
check=True
)
import os
import spaces
import gradio as gr
import cv2
import numpy as np
import PIL.Image
import torch
# Global detector (initialized on first GPU call)
detector = None
@spaces.GPU
def detect(image, face_score_threshold, landmark_score_threshold):
"""Detect anime faces and landmarks in the image."""
global detector
# Lazy initialization with auto device detection
if detector is None:
from anime_face_detector import create_detector
# Auto-detect device: use GPU if available, otherwise CPU
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')
detector = create_detector('yolov3', device=device)
# Convert RGB to BGR for OpenCV
image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# Run detection
preds = detector(image_bgr)
# Draw results
res = image_bgr.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)
# Line thickness based on face size
lt = max(2, int(3 * (box[2:] - box[:2]).max() / 256))
# Draw bounding box
cv2.rectangle(res, tuple(box[:2]), tuple(box[2:]), (0, 255, 0), lt)
# Draw confidence score
cv2.putText(res, f'{score * 100:.1f}%', (box[0], box[1] - 5),
cv2.FONT_HERSHEY_SIMPLEX, lt / 3, (255, 255, 255), thickness=max(lt // 2, 1))
# Draw keypoints
pred_pts = pred['keypoints']
for *pt, kpt_score in pred_pts:
if kpt_score < landmark_score_threshold:
color = (0, 255, 255) # Yellow for low confidence
else:
color = (0, 0, 255) # Red for high confidence
pt = tuple(np.round(pt).astype(int))
cv2.circle(res, pt, lt, color, cv2.FILLED)
# Convert BGR to RGB for output
res = cv2.cvtColor(res, cv2.COLOR_BGR2RGB)
return PIL.Image.fromarray(res)
# Download sample image
def download_sample():
sample_path = 'input.jpg'
if not os.path.exists(sample_path):
torch.hub.download_url_to_file(
'https://raw.githubusercontent.com/hysts/anime-face-detector/main/assets/input.jpg',
sample_path
)
return sample_path
# Create Gradio interface
sample_path = download_sample()
demo = gr.Interface(
fn=detect,
inputs=[
gr.Image(type='numpy', label='Input Image'),
gr.Slider(0, 1, step=0.05, value=0.5, label='Face Score Threshold'),
gr.Slider(0, 1, step=0.05, value=0.3, label='Landmark Score Threshold'),
],
outputs=gr.Image(type='pil', label='Output'),
title='Anime Face Detector (GPU)',
description='Detect anime faces and 28 facial landmarks using YOLOv3 + HRNetV2. Powered by Zero GPU.',
article='<a href="https://github.com/hysts/anime-face-detector">GitHub</a> | <a href="https://github.com/ayutaz/anime-face-detector">Fork with GPU Docker</a>',
examples=[
[sample_path, 0.5, 0.3],
],
cache_examples=False,
)
if __name__ == '__main__':
demo.launch()