Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,638 Bytes
3d5837a |
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 39 40 41 42 43 44 45 46 |
import os
import folder_paths
import impact.core as core
import impact.subcore as subcore
from impact.utils import add_folder_path_and_extensions
version_code = 22
print(f"### Loading: ComfyUI-Impact-Pack (Subpack: V0.6)")
model_path = folder_paths.models_dir
add_folder_path_and_extensions("ultralytics_bbox", [os.path.join(model_path, "ultralytics", "bbox")], folder_paths.supported_pt_extensions)
add_folder_path_and_extensions("ultralytics_segm", [os.path.join(model_path, "ultralytics", "segm")], folder_paths.supported_pt_extensions)
add_folder_path_and_extensions("ultralytics", [os.path.join(model_path, "ultralytics")], folder_paths.supported_pt_extensions)
class UltralyticsDetectorProvider:
@classmethod
def INPUT_TYPES(s):
bboxs = ["bbox/"+x for x in folder_paths.get_filename_list("ultralytics_bbox")]
segms = ["segm/"+x for x in folder_paths.get_filename_list("ultralytics_segm")]
return {"required": {"model_name": (bboxs + segms, )}}
RETURN_TYPES = ("BBOX_DETECTOR", "SEGM_DETECTOR")
FUNCTION = "doit"
CATEGORY = "ImpactPack"
def doit(self, model_name):
model_path = folder_paths.get_full_path("ultralytics", model_name)
model = subcore.load_yolo(model_path)
if model_name.startswith("bbox"):
return subcore.UltraBBoxDetector(model), core.NO_SEGM_DETECTOR()
else:
return subcore.UltraBBoxDetector(model), subcore.UltraSegmDetector(model)
NODE_CLASS_MAPPINGS = {
"UltralyticsDetectorProvider": UltralyticsDetectorProvider
}
NODE_DISPLAY_NAME_MAPPINGS = {
}
|