| import os |
| import folder_paths |
| from . import subcore |
| from . import utils |
| import logging |
|
|
|
|
| def update_model_paths(model_path): |
| utils.add_folder_path_and_extensions("ultralytics_bbox", [os.path.join(model_path, "ultralytics", "bbox")], folder_paths.supported_pt_extensions) |
| utils.add_folder_path_and_extensions("ultralytics_segm", [os.path.join(model_path, "ultralytics", "segm")], folder_paths.supported_pt_extensions) |
| utils.add_folder_path_and_extensions("ultralytics", [os.path.join(model_path, "ultralytics")], folder_paths.supported_pt_extensions) |
| logging.info(f'[Impact Subpack] ultralytics_bbox: {", ".join(folder_paths.folder_names_and_paths["ultralytics_bbox"][0])}') |
| logging.info(f'[Impact Subpack] ultralytics_segm: {", ".join(folder_paths.folder_names_and_paths["ultralytics_segm"][0])}') |
|
|
| update_model_paths(folder_paths.models_dir) |
| if 'download_model_base' in folder_paths.folder_names_and_paths: |
| update_model_paths(folder_paths.get_folder_paths('download_model_base')[0]) |
|
|
|
|
| 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) |
|
|
| if model_path is None: |
| if model_name.startswith('bbox/'): |
| model_path = folder_paths.get_full_path("ultralytics_bbox", model_name[5:]) |
| elif model_name.startswith('segm/'): |
| model_path = folder_paths.get_full_path("ultralytics_segm", model_name[5:]) |
|
|
| if model_path is None: |
| logging.error(f"[Impact Subpack] model file '{model_name}' is not found in one of the following directories:") |
|
|
| cands = [] |
| cands.extend(folder_paths.get_folder_paths("ultralytics")) |
| if model_name.startswith('bbox/'): |
| cands.extend(folder_paths.get_folder_paths("ultralytics_bbox")) |
| elif model_name.startswith('segm/'): |
| cands.extend(folder_paths.get_folder_paths("ultralytics_segm")) |
|
|
| formatted_cands = "\n\t".join(cands) |
| logging.error(f'\t{formatted_cands}\n') |
|
|
| raise ValueError(f"[Impact Subpack] model file '{model_name}' is not found.") |
|
|
| model = subcore.load_yolo(model_path) |
|
|
| if model_name.startswith("bbox"): |
| return subcore.UltraBBoxDetector(model), subcore.NO_SEGM_DETECTOR() |
| else: |
| return subcore.UltraBBoxDetector(model), subcore.UltraSegmDetector(model) |
|
|
|
|
| NODE_CLASS_MAPPINGS = { |
| "UltralyticsDetectorProvider": UltralyticsDetectorProvider |
| } |
|
|
|
|
| NODE_DISPLAY_NAME_MAPPINGS = { |
|
|
| } |
|
|
