49 lines
1.2 KiB
Python
49 lines
1.2 KiB
Python
|
|
import comfy.model_management as model_management
|
|
|
|
from .canny import CannyDetector
|
|
from .hed import HEDdetector
|
|
from .leres import LeresDetector
|
|
from .utils import annotator_wrapper
|
|
|
|
|
|
def canny(image, low_threshold, high_threshold):
|
|
annotator = CannyDetector()
|
|
|
|
low_threshold = int(low_threshold * 255)
|
|
high_threshold = int(high_threshold * 255)
|
|
|
|
def annotate(np_image):
|
|
return annotator(np_image, low_threshold=low_threshold, high_threshold=high_threshold)
|
|
|
|
return annotator_wrapper(image, annotate)
|
|
|
|
|
|
def leres(image, rm_nearest, rm_background, annotator_model):
|
|
annotator = LeresDetector.from_pretrained(annotator_model).to(model_management.get_torch_device())
|
|
|
|
rm_nearest = rm_nearest * 100
|
|
rm_background = (1.0 - rm_background) * 100
|
|
|
|
def annotate(np_image):
|
|
return annotator(np_image, thr_a=rm_background, thr_b=rm_nearest)
|
|
|
|
out = annotator_wrapper(image, annotate)
|
|
|
|
del annotator
|
|
|
|
return out
|
|
|
|
|
|
def hed(image, annotator_model):
|
|
annotator = HEDdetector.from_pretrained(annotator_model).to(model_management.get_torch_device())
|
|
|
|
def annotate(np_image):
|
|
return annotator(np_image, safe=True)
|
|
|
|
out = annotator_wrapper(image, annotate)
|
|
|
|
del annotator
|
|
|
|
return out
|