Spaces:
Running
Running
| import os | |
| import logging | |
| from pathlib import Path | |
| from dotenv import load_dotenv | |
| from ultralytics import YOLO | |
| load_dotenv(override=True) | |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| # ---------------- LOGGER ---------------- | |
| LOG_FILE = os.path.join(BASE_DIR, "api.log") | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s [%(levelname)s] %(message)s", | |
| handlers=[logging.FileHandler(LOG_FILE), logging.StreamHandler()] | |
| ) | |
| logger = logging.getLogger("WildlifeLogger") | |
| ENV = os.getenv("ENV", "DEV").upper() | |
| logger.info(f"Running in {ENV}") | |
| UPLOAD_DIR = os.getenv("HF_BUCKET_ID") | |
| if not UPLOAD_DIR: | |
| raise ValueError("HF_BUCKET_ID is not set in environment variables") | |
| # ---------------- UPLOAD RULES --------w-------- | |
| MIN_IMAGES = 1 | |
| MAX_IMAGES = 1000 | |
| ALLOWED_EXTENSIONS = {"jpg", "jpeg", "png", "webp"} | |
| # ---------------- YOLO MODELS ---------------- | |
| try: | |
| logger.info("Loading YOLO models...") | |
| DETECT_MODEL = YOLO("api/walidlife_models/detect/deer.onnx") | |
| BUCK_DOE_MODEL = YOLO("api/walidlife_models/classify/Buck_classificationt.onnx", task="classify") | |
| BUCK_TYPE_MODEL = YOLO("api/walidlife_models/classify/mules_vs_whitetails.onnx", task="classify") | |
| logger.info("YOLO models loaded") | |
| except Exception as e: | |
| logger.error(f"YOLO load failed: {e}") | |
| DETECT_MODEL = BUCK_DOE_MODEL = BUCK_TYPE_MODEL = None |