diff --git a/config.py b/config.py index f22d374..984e35e 100644 --- a/config.py +++ b/config.py @@ -1,3 +1,32 @@ +# YOLOX configurations YOLOX_CONFIG_FILE = "modules/yolox/config/yolox_x_8x8_300e_coco-EARS-white.py" YOLOX_CHECKPOINT_FILE = "models/yolox/yolox_x_White_only.pth" DEVICE = "cuda:0" + +# Colors for different models (R,G,B format) +CONV_COLOR = (0, 255, 0) # Green +XGBOOST_COLOR = (255, 0, 0) # Red +LIGHTGBM_COLOR = (0, 0, 255) # Blue +RESNET_COLOR = (255, 165, 0) # Orange +EFFICIENTNET_COLOR = (0, 0, 255) # Blue +MOBILENET_COLOR = (255, 0, 0) # Red +SQUEEZENET_COLOR = (128, 0, 128) # Purple + +# Model execution settings +CONV_ENABLED = True +XGBOOST_ENABLED = False +LIGHTGBM_ENABLED = False +NORMALIZE_ENABLED = False +POSENET_ENABLED = True +RTMPOSE_ENABLED = False +MOBILENETV1SSD_ENABLED = True +YOLOX_ENABLED = False + +# Neural network model settings +RESNET_ENABLED = False +EFFICIENTNET_ENABLED = False +MOBILENET_ENABLED = False +SQUEEZENET_ENABLED = False + +# Video configuration +VIDEO_PATH = "./video/Test1.mp4" diff --git a/main.py b/main.py index c20442b..c7993bf 100644 --- a/main.py +++ b/main.py @@ -7,7 +7,6 @@ import cv2 import numpy as np import pandas as pd -from dotenv import load_dotenv from mmdet.apis import DetInferencer, inference_detector, init_detector # New imports for RTMPose @@ -22,31 +21,22 @@ from util.calc_ste_position import CalcStethoscopePosition from util.ears_ai import EarsAI -# Load environment variables -load_dotenv() - -# Get colors from environment variables -CONV_COLOR = tuple( - map(int, os.getenv("CONV_COLOR", "0,255,0").split(",")) -) # Default: Green -XGBOOST_COLOR = tuple( - map(int, os.getenv("XGBOOST_COLOR", "255,0,0").split(",")) -) # Default: Red -LIGHTGBM_COLOR = tuple( - map(int, os.getenv("LIGHTGBM_COLOR", "0,0,255").split(",")) -) # Default: Blue +# Get colors from config +CONV_COLOR = config.CONV_COLOR +XGBOOST_COLOR = config.XGBOOST_COLOR +LIGHTGBM_COLOR = config.LIGHTGBM_COLOR # Get model execution settings -CONV_ENABLED = os.getenv("CONV_ENABLED", "True").lower() == "true" -XGBOOST_ENABLED = os.getenv("XGBOOST_ENABLED", "True").lower() == "true" -LIGHTGBM_ENABLED = os.getenv("LIGHTGBM_ENABLED", "True").lower() == "true" -POSENET_ENABLED = os.getenv("PoseNet_ENABLED", "True").lower() == "true" -RTMPOSE_ENABLED = os.getenv("RTMPose_ENABLED", "False").lower() == "true" -MobileNetV1SSD_ENABLED = os.getenv("MobileNetV1SSD_ENABLED", "False").lower() == "true" -YOLOX_ENABLED = os.getenv("YOLOX_ENABLED", "False").lower() == "true" +CONV_ENABLED = config.CONV_ENABLED +XGBOOST_ENABLED = config.XGBOOST_ENABLED +LIGHTGBM_ENABLED = config.LIGHTGBM_ENABLED +POSENET_ENABLED = config.POSENET_ENABLED +RTMPOSE_ENABLED = config.RTMPOSE_ENABLED +MobileNetV1SSD_ENABLED = config.MOBILENETV1SSD_ENABLED +YOLOX_ENABLED = config.YOLOX_ENABLED # Get normalization setting -NORMALIZE_ENABLED = os.getenv("NORMALIZE_ENABLED", "False").lower() == "true" +NORMALIZE_ENABLED = config.NORMALIZE_ENABLED def init_yolox():