CodeProject.AI-ObjectDetectionYOLOv5-3.1/options.py at main · codeproject/CodeProject.AI-ObjectDetectionYOLOv5-3.1 · GitHub

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
import os
from codeproject_ai_sdk import ModuleOptions

class Settings:

    def __init__(self, RESOLUTION, STD_MODEL_NAME):
        self.RESOLUTION     = RESOLUTION
        self.STD_MODEL_NAME = STD_MODEL_NAME


class Options:

    def __init__(self):
        # -------------------------------------------------------------------------
        # Setup constants

        # see https://github.com/ultralytics/yolov5 for resolution data
        self.MODEL_SETTINGS = {
            # "tiny":   Settings(STD_MODEL_NAME = "yolov5n", RESOLUTION = 256), # 640
            "small":  Settings(STD_MODEL_NAME = "yolov5s", RESOLUTION = 256),
            "medium": Settings(STD_MODEL_NAME = "yolov5m", RESOLUTION = 416),
            "large":  Settings(STD_MODEL_NAME = "yolov5l", RESOLUTION = 640),
            # "huge":   Settings(STD_MODEL_NAME = "yolov5x", RESOLUTION = 640)  # Not yet included
        }

        # -------------------------------------------------------------------------
        # Setup values

        self._show_env_variables = True

        self.app_dir            = os.path.normpath(ModuleOptions.getEnvVariable("APPDIR", os.getcwd()))
        self.models_dir         = os.path.normpath(ModuleOptions.getEnvVariable("MODELS_DIR", f"{self.app_dir}/assets"))
        self.custom_models_dir  = os.path.normpath(ModuleOptions.getEnvVariable("CUSTOM_MODELS_DIR", f"{self.app_dir}/custom-models"))

        self.sleep_time         = 0.01

        self.model_size         = ModuleOptions.getEnvVariable("MODEL_SIZE", "Medium")   # small, medium, large //, nano, x-large
        self.use_CUDA           = ModuleOptions.getEnvVariable("USE_CUDA",   "True")     # True / False
        self.use_MPS            = False   # Default is False, but we'll enable if possible

        # Normalise input
        self.model_size         = self.model_size.lower()
        self.use_CUDA           = ModuleOptions.enable_GPU and self.use_CUDA.lower() == "true"

        if self.model_size not in [ "small", "medium", "large" ]:
            self.model_size = "small"

        # Get settings
        settings = self.MODEL_SETTINGS[self.model_size]
        self.resolution_pixels = settings.RESOLUTION
        self.std_model_name    = settings.STD_MODEL_NAME

        # -------------------------------------------------------------------------
        # dump the important variables

        if self._show_env_variables:
            print(f"Debug: APPDIR:      {self.app_dir}")
            print(f"Debug: MODEL_SIZE:  {self.model_size}")
            print(f"Debug: MODELS_DIR:  {self.models_dir}")