bytetrack-counter-go/pkg/config/config.go

package config

import (
	"os"
	"strconv"
	"strings"
)

type Config struct {
	OutputDir   string
	DBPath      string
	StateFile   string
	Source      string
	ModelPath   string
	CameraName  string
	ObjectLabel string
	ClassAyam   string
	ClassTalenan string

	LineX          *int
	LineXFrac      float64
	CrossDirection string

	ImgSz        int
	Half         bool
	Conf         float64
	CoreMask     int
	NumClasses   int
	ScoreSigmoid bool

	TrackHighThresh float64
	TrackLowThresh  float64
	TrackMatchThresh float64
	TrackBuffer    int
	TrackMinHits   int

	DailyCutoffTime            string
	BatchTimeoutSec            float64
	IgnoreBatchLabelTimeoutSec float64
	MinObjectPerBatch          int
	MinDurationPerBatch        int

	ExportCSV   bool
	CrossCSV    string
	RecordVideo bool

	WarmupFrames         int
	ReconnectDelaySec    int
	MaxReconnectAttempts int
	FlushEveryNFrames    int
	TrackedPruneSec      float64
	VideoSegmentSec      int
	OutputFPS            int

	LiveStreamEnabled     bool
	LiveStreamFramePath   string
	LiveStreamQuality     int
	LiveStreamEveryN      int
	RTSPFFmpegOptions     string

	IsLive bool
}

func Load() *Config {
	cfg := &Config{
		OutputDir:   getEnv("OUTPUT_DIR", "/opt/jetson-counter"),
		DBPath:      getEnv("DB_PATH", ""),
		StateFile:   getEnv("STATE_FILE", ""),
		Source:      getEnv("SOURCE", "rtsp://user:pass@192.168.0.100:554/stream1"),
		ModelPath:   getEnv("MODEL_PATH", "/opt/jetson-counter/yolo11n.rknn"),
		CameraName:  getEnv("CAMERA_NAME", "CC1"),
		ObjectLabel: getEnv("OBJECT_LABEL", "ayam-potong"),
		ClassAyam:   getEnv("CLASS_AYAM", "ayam"),
		ClassTalenan: getEnv("CLASS_TALENAN", "talenan"),

		LineXFrac:      getEnvFloat("LINE_X_FRAC", 0.5),
		CrossDirection: strings.ToLower(getEnv("CROSS_DIRECTION", "rtl")),

		ImgSz:        getEnvInt("IMGSZ", 320),
		Half:         getEnvBool("HALF"),
		Conf:         getEnvFloat("CONF", 0.3),
		CoreMask:     getEnvInt("CORE_MASK", 1),
		NumClasses:   getEnvInt("NUM_CLASSES", 2),
		ScoreSigmoid: getEnvBool("SCORE_SIGMOID"),

		TrackHighThresh:  getEnvFloat("TRACK_HIGH_THRESH", 0.5),
		TrackLowThresh:   getEnvFloat("TRACK_LOW_THRESH", 0.1),
		TrackMatchThresh: getEnvFloat("TRACK_MATCH_THRESH", 0.8),
		TrackBuffer:      getEnvInt("TRACK_BUFFER", 30),
		TrackMinHits:     getEnvInt("TRACK_MIN_HITS", 3),

		DailyCutoffTime:            getEnv("DAILY_CUTOFF_TIME", "20:00"),
		BatchTimeoutSec:            getEnvFloat("BATCH_TIMEOUT_SECONDS", 300),
		IgnoreBatchLabelTimeoutSec: getEnvFloat("IGNORE_BATCH_LABEL_TIMEOUT_SECONDS", 30),
		MinObjectPerBatch:          getEnvInt("MIN_OBJECT_PER_BATCH", 60),
		MinDurationPerBatch:        getEnvInt("MIN_DURATION_PER_BATCH", 60),

		ExportCSV:   getEnvBool("EXPORT_CSV"),
		CrossCSV:    getEnv("CROSS_CSV", ""),
		RecordVideo: getEnvBool("RECORD_VIDEO"),

		WarmupFrames:         getEnvInt("WARMUP_FRAMES", 30),
		ReconnectDelaySec:    getEnvInt("RECONNECT_DELAY_SEC", 3),
		MaxReconnectAttempts: getEnvInt("MAX_RECONNECT_ATTEMPTS", 0),
		FlushEveryNFrames:    getEnvInt("FLUSH_EVERY_N_FRAMES", 100),
		TrackedPruneSec:      getEnvFloat("TRACKED_PRUNE_SEC", 300),
		VideoSegmentSec:      getEnvInt("VIDEO_SEGMENT_SEC", 3600),
		OutputFPS:            getEnvInt("OUTPUT_FPS", 15),

		LiveStreamEnabled:   getEnvBool("LIVE_STREAM_ENABLED"),
		LiveStreamFramePath: getEnv("LIVE_STREAM_FRAME_PATH", "/dev/shm/jetson-counter/live_frame.jpg"),
		LiveStreamQuality:   getEnvInt("LIVE_STREAM_QUALITY", 75),
		LiveStreamEveryN:    getEnvInt("LIVE_STREAM_EVERY_N", 2),
		RTSPFFmpegOptions:   getEnv("OPENCV_FFMPEG_CAPTURE_OPTIONS", "rtsp_transport;tcp|fflags;nobuffer|flags;low_delay"),
	}

	if cfg.DBPath == "" {
		cfg.DBPath = cfg.OutputDir + "/jetson_counter.db"
	}
	if cfg.StateFile == "" {
		cfg.StateFile = cfg.OutputDir + "/current_batch.json"
	}
	if cfg.CrossCSV == "" {
		cfg.CrossCSV = cfg.OutputDir + "/batch_crossings.csv"
	}

	if lx := os.Getenv("LINE_X"); lx != "" {
		v, err := strconv.Atoi(lx)
		if err == nil {
			cfg.LineX = &v
		}
	}

	cfg.IsLive = strings.HasPrefix(strings.ToLower(cfg.Source), "rtsp://") ||
		strings.HasPrefix(strings.ToLower(cfg.Source), "http://")

	return cfg
}

func getEnv(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

func getEnvInt(key string, def int) int {
	s := os.Getenv(key)
	if s == "" {
		return def
	}
	v, err := strconv.Atoi(s)
	if err != nil {
		return def
	}
	return v
}

func getEnvFloat(key string, def float64) float64 {
	s := os.Getenv(key)
	if s == "" {
		return def
	}
	v, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return def
	}
	return v
}

func getEnvBool(key string) bool {
	return strings.ToLower(os.Getenv(key)) == "true"
}