constants.py

import os

from dotenv import load_dotenv

load_dotenv()

# Model checkpoints and repo names
CLAS_FILENAME = os.getenv("CLAS_FILENAME")
DET_FILENAME = os.getenv("DET_FILENAME")
HF_CLAS_REPO_NAME = os.getenv("HF_CLAS_REPO_NAME")
HF_DET_REPO_NAME = os.getenv("HF_DET_REPO_NAME")

# Models paths
MODELS_PATH = "models"
CLAS_FILEPATH = f"{MODELS_PATH}/{CLAS_FILENAME}"
DET_FILEPATH = f"{MODELS_PATH}/{DET_FILENAME}"

# Other constants
CLASSES = ["cardboard", "compost", "glass", "metal", "paper", "plastic", "trash"]
CLAS_THRESHOLD = 0.5
DET_NAME = "tf_efficientdet_d2"
DET_THRESHOLD = 0.17
DEVICE = "cpu"
OUTPUT_IMG_FILEPATH = "classified_image.jpg"