diff --git a/.gitignore b/.gitignore index 464de57..44462f4 100644 --- a/.gitignore +++ b/.gitignore @@ -178,4 +178,4 @@ models/ !modules/posenet/models output/ -data/ \ No newline at end of file +data/ diff --git a/config.py b/config.py index 9de56cc..44850ea 100644 --- a/config.py +++ b/config.py @@ -4,7 +4,7 @@ DEVICE = "cuda:0" # Colors for different models (R,G,B format) -CONV_COLOR = (255, 165, 0) # Green +CONV_COLOR = (0, 255, 0) # Green XGBOOST_COLOR = (255, 0, 0) # Red LIGHTGBM_COLOR = (0, 0, 255) # Blue RESNET_COLOR = (255, 165, 0) # Orange @@ -15,7 +15,7 @@ # Model execution settings CONV_ENABLED = True XGBOOST_ENABLED = False -LIGHTGBM_ENABLED = False +LIGHTGBM_ENABLED = True NORMALIZE_ENABLED = False POSENET_ENABLED = False RTMPOSE_ENABLED = True diff --git a/main.py b/main.py index db0b7ff..4d4fed1 100644 --- a/main.py +++ b/main.py @@ -470,10 +470,10 @@ if LIGHTGBM_ENABLED: lgb_model_x = load_model( - "./models/lgb_stethoscope_calc_x_best_model-Fold4.pkl" + "./models/LightGBM/stethoscope_calc_x_best_model.pkl" ) lgb_model_y = load_model( - "./models/lgb_stethoscope_calc_y_best_model-Fold4.pkl" + "./models/LightGBM/stethoscope_calc_y_best_model.pkl" ) if XGBOOST_ENABLED: xg_model_x = load_model( diff --git a/modules/posenet/models/model_factory.py b/modules/posenet/models/model_factory.py index 9051492..e0cf322 100644 --- a/modules/posenet/models/model_factory.py +++ b/modules/posenet/models/model_factory.py @@ -1,10 +1,10 @@ -import torch import os +import torch -from modules.posenet.models.mobilenet_v1 import MobileNetV1, MOBILENET_V1_CHECKPOINTS +from modules.posenet.models.mobilenet_v1 import MOBILENET_V1_CHECKPOINTS, MobileNetV1 -MODEL_DIR = "./models" +MODEL_DIR = "./models/posenet" DEBUG_OUTPUT = False diff --git a/modules/util/const.py b/modules/util/const.py index dc9341f..36bac9c 100644 --- a/modules/util/const.py +++ b/modules/util/const.py @@ -9,8 +9,8 @@ BAUDRATE = 115200 # SSDモデルのセットアップ -MODEL_PATH = "./models/mb1-ssd-second.pth" -LABEL_PATH = "./models/voc-model-labels.txt" +MODEL_PATH = "./models/ssd/mb1-ssd-second.pth" +LABEL_PATH = "./models/ssd/voc-model-labels.txt" # EARS音源のセットアップ EARS_MAP_PATH = "img/map/"