Files
HOSPITAL_CCTV/demo_const.py
2026-02-25 15:05:58 +09:00

128 lines
3.8 KiB
Python

import project_config
from config_loader import CFG, get_weights_dir
from color_table import dev_color_table, demo_color_table
# === 카메라 프리셋 ===
_cameras = CFG.get('cameras', {})
MG = _cameras.get('MG', '')
LF = _cameras.get('LF', '')
FC = _cameras.get('FC', '')
MG_74_LTE = _cameras.get('MG_74_LTE', '')
FERMAT = _cameras.get('FERMAT', '')
# === 소스 ===
_source_raw = CFG.get('source', '')
SOURCE = _cameras.get(_source_raw, _source_raw) # 프리셋 이름이면 URL로, 아니면 그대로
SAVE_PATH = CFG.get('save_path', './250630_result')
# === 모델 가중치 경로 ===
WEIGHTS_PATH = get_weights_dir()
_weights = CFG.get('weights', {})
WEIGHTS_POSE = WEIGHTS_PATH + "/" + _weights.get('pose', 'yolov8l-pose.pt')
WEIGHTS_YOLO_HELMET = WEIGHTS_PATH + "/" + _weights.get('helmet', 'yolov8_dev1_97.pt')
# PT
HELMET_KPT = [0, 1, 2]
RUBBER_INSULATED_SLEEVE_KPT = [7, 8]
SUIT_TOP_KPT = [6, 5]
SUIT_BOTTOM_KPT = [13, 14]
HELMET_CLASS_NAME = ['safety_helmet_off', 'safety_helmet_on']
# PT dev
if project_config.PT_TYPE == 'dev':
color_table = dev_color_table
WEIGHTS_YOLO = WEIGHTS_PATH + "/" + _weights.get('od_dev', 'yolov11m_dev1_8.pt')
UNVISIBLE_CLS = [4, 5, 6, 7, 8, 9]
OFF_CLASS_LIST = [2, 4, 12, 14]
OFF_TRIGGER_CLASS_LIST = [2, 4] # (helmet, gloves)
PPE_CLASS_LIST = list(range(2, 19))
HELMET_CID = [2, 3] # OFF, ON
RUBBER_INSULATED_SLEEVE_CID = [10, 11] # OFF, ON
SUIT_TOP_CID = [] # OFF, ON
SUIT_BOTTOM_CID = [] # OFF, ON
HELMET_ON_CID = 3
GLOVES_WORK_ON_CID = 6
BOOTS_ON_CID = 13
TRAFFIC_CONE_CID = 21
LIGHT_STICK_CLS_ID = 23
# PT DEMO
elif project_config.PT_TYPE == 'demo':
color_table = demo_color_table
WEIGHTS_YOLO = WEIGHTS_PATH + "/" + _weights.get('od_demo', 'yolov8_dev1_66.pt')
UNVISIBLE_CLS = [3, 4, 5]
OFF_CLASS_LIST = [1, 3, 6, 8]
OFF_TRIGGER_CLASS_LIST = [1, 3] # (helmet, gloves)
PPE_CLASS_LIST = list(range(1, 9))
HELMET_CID = [1, 2] # OFF, ON
RUBBER_INSULATED_SLEEVE_CID = []
SUIT_TOP_CID = [] # OFF, ON
SUIT_BOTTOM_CID = [6, 7] # OFF, ON
HELMET_ON_CID = 2
GLOVES_WORK_ON_CID = 4
BOOTS_ON_CID = 7
TRAFFIC_CONE_CID = 11
# === 모델 파라미터 ===
_display = CFG.get('display', {})
HPE_FRAME_CHECK_MAX_COUNT = CFG.get('hpe', {}).get('hpe_frame_check_max_count', 3)
PPE_UNION_MIN_PERCENT = _display.get('ppe_union_min_percent', 0.9)
MODEL_CONFIDENCE = CFG.get('model_confidence', 0.5)
MODEL_IMAGE_SIZE = CFG.get('model_image_size', 640)
WD_BORDER_COLOR = [0, 255, 255]
WD_THICKNESS_ON = False
BORDER_OD_TEXT = "Danger(OD)"
BORDER_HPE_TEXT = "Danger(HPE)"
KPT_MIN_CONFIDENCE = CFG.get('kpt_min_confidence', 0.5)
# loadstreams
LOADSTREAMS_IMG_BUFFER = _display.get('loadstreams_img_buffer', 5)
_fhd = _display.get('fhd_resolution', [1920, 1080])
FHD_RESOLUTION = tuple(_fhd)
BORDER_THICKNESS = _display.get('border_thickness', 40)
BORDER_THICKNESS_HALF = _display.get('border_thickness_half', 20)
NORMAL_THICKNESS = _display.get('normal_thickness', 2)
WARNING_THICKNESS = _display.get('warning_thickness', 4)
HPE_THICKNESS_RAITO = _display.get('hpe_thickness_ratio', 1)
TEXT_SIZE = _display.get('text_size', 1)
TEXT_THICKNESS = _display.get('text_thickness', 0)
TEXT_OD_STARTING_POINT = (20, 30)
TEXT_HPE_STARTING_POINT = (20, 30)
# POSE HUMAN LABEL COLOR (BGR)
POSE_NORMAL_COLOR = [0, 0, 0]
POSE_CROSS_COLOR = [0, 0, 255]
POSE_FALL_COLOR = [0, 0, 255]
POSE_ANGLE_ARM_COLOR = [0, 51, 102]
TEXT_COLOR_WHITE = [255, 255, 255]
TEXT_COLOR_BLACK = [0, 0, 0]
text_color_white_list = [
"truck",
"safety_gloves_insulated_on_1",
"safety_gloves_insulated_on_2",
"safety_rubber_insulated_sleeve_on",
"safety_belt_swing_on",
"safety_vest_on",
"safety_suit_top_on",
"sign_board_traffic"
]