File size: 1,570 Bytes
9dce458 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
import numpy as np
from .default import DefaultDetector
from .dbnet_convnext import DBConvNextDetector
from .ctd import ComicTextDetector
from .craft import CRAFTDetector
from .none import NoneDetector
from .common import CommonDetector, OfflineDetector
DETECTORS = {
'default': DefaultDetector,
'dbconvnext': DBConvNextDetector,
'ctd': ComicTextDetector,
'craft': CRAFTDetector,
'none': NoneDetector,
}
detector_cache = {}
def get_detector(key: str, *args, **kwargs) -> CommonDetector:
if key not in DETECTORS:
raise ValueError(f'Could not find detector for: "{key}". Choose from the following: %s' % ','.join(DETECTORS))
if not detector_cache.get(key):
detector = DETECTORS[key]
detector_cache[key] = detector(*args, **kwargs)
return detector_cache[key]
async def prepare(detector_key: str):
detector = get_detector(detector_key)
if isinstance(detector, OfflineDetector):
await detector.download()
async def dispatch(detector_key: str, image: np.ndarray, detect_size: int, text_threshold: float, box_threshold: float, unclip_ratio: float,
invert: bool, gamma_correct: bool, rotate: bool, auto_rotate: bool = False, device: str = 'cpu', verbose: bool = False):
detector = get_detector(detector_key)
if isinstance(detector, OfflineDetector):
await detector.load(device)
return await detector.detect(image, detect_size, text_threshold, box_threshold, unclip_ratio, invert, gamma_correct, rotate, auto_rotate, verbose)
|