setup_env.py

"""
setup_env.py — One-click environment setup for Face Mask Detection.

Downloads the required AI model files into the ``models/`` directory.
Safe to re-run; files already present are skipped automatically.

Usage
-----
    python setup_env.py
"""

import os
import sys
import logging
import warnings

warnings.filterwarnings("ignore")
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

import requests

from src import config

# ---------------------------------------------------------------------------
# Logging — must be configured before any log calls are made
# ---------------------------------------------------------------------------
logging.basicConfig(format=config.LOG_FORMAT, level=logging.INFO)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Remote model URLs
# ---------------------------------------------------------------------------
YUNET_URL = (
    "https://github.com/opencv/opencv_zoo/raw/main/models/"
    "face_detection_yunet/face_detection_yunet_2023mar.onnx"
)

MASK_MODEL_URL = "https://github.com/chandrikadeb7/Face-Mask-Detection/raw/master/mask_detector.model"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def download_file(url: str, filepath: str) -> bool:
    """
    Download *url* to *filepath* with a simple byte-count progress indicator.

    Returns ``True`` on success, ``False`` on error.
    """
    filename = os.path.basename(filepath)
    logger.info(f"Downloading {filename} ...")
    try:
        response = requests.get(url, stream=True, timeout=60)
        response.raise_for_status()

        total = int(response.headers.get("Content-Length", 0))
        downloaded = 0

        with open(filepath, "wb") as fh:
            for chunk in response.iter_content(chunk_size=65_536):
                fh.write(chunk)
                downloaded += len(chunk)
                if total:
                    pct = downloaded / total * 100
                    # Overwrite the same terminal line for a progress effect
                    print(
                        f"\r  {filename}: {downloaded / 1_048_576:.1f} MB / "
                        f"{total / 1_048_576:.1f} MB  ({pct:.0f}%)",
                        end="",
                        flush=True,
                    )

        print()  # newline after progress line
        size_mb = os.path.getsize(filepath) / 1_048_576
        logger.info(f"Saved {filename} ({size_mb:.2f} MB) → {filepath}")
        return True

    except requests.exceptions.RequestException as exc:
        logger.error(f"Network error downloading {filename}: {exc}")
        # Remove incomplete file to avoid loading a corrupt artefact
        if os.path.exists(filepath):
            os.remove(filepath)
        return False
    except OSError as exc:
        logger.error(f"File write error for {filename}: {exc}")
        return False


def ensure_directory(path: str) -> None:
    """Create *path* (and any missing parents) if it does not already exist."""
    os.makedirs(path, exist_ok=True)
    logger.debug(f"Directory ready: {path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    print("=" * 50)
    print("  Face Mask Detection — Environment Setup")
    print("=" * 50)

    # 1. Ensure models directory exists
    ensure_directory(config.MODELS_DIR)

    # 2. YuNet face detector
    if not os.path.exists(config.FACE_MODEL_YUNET):
        success = download_file(YUNET_URL, config.FACE_MODEL_YUNET)
        if not success:
            logger.error("YuNet download failed. Face detection will not work.")
    else:
        logger.info(f"Found YuNet model: {config.FACE_MODEL_YUNET}")

    # 3. Mask classifier (saved as .h5 regardless of remote file extension)
    if not os.path.exists(config.MASK_MODEL_PATH):
        logger.info("Downloading pre-trained mask classifier...")
        success = download_file(MASK_MODEL_URL, config.MASK_MODEL_PATH)
        if not success:
            logger.error(
                "Mask model download failed.\n"
                "You can manually download it from:\n"
                f"  {MASK_MODEL_URL}\n"
                f"and save it to: {config.MASK_MODEL_PATH}"
            )
    else:
        logger.info(f"Found mask model: {config.MASK_MODEL_PATH}")

    print()
    print("Setup complete. Run the application with:")
    print("  python main.py")
    print()


if __name__ == "__main__":
    main()