Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
subpack_nodes.py 2.88 KB
Newer Older
wuxk1's avatar
optim for test1,2,5  
wuxk1 committed
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70 
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 = {

}