app.py

from fastapi import FastAPI
from prediction.controller.endpoints import router
from auth.controller.endpoints import router as auth_router
from database import Base, engine

import os

# Import the artifact seeder
from artifact_seeder import download_models

pet_prediction_app = FastAPI()
Base.metadata.create_all(bind=engine)

# Download models if they don't exist
artifacts_dir = os.path.join("prediction", "artifacts")
required_models = [
    "resnet50_imagenet.pth",
    "deeplabv3_resnet50.pth",
    "yolov5s.pt"
]

models_exist = all(
    os.path.exists(os.path.join(artifacts_dir, model))
    for model in required_models
)

if not models_exist:
    print("Some models missing, downloading...")
    download_models()
else:
    print("All models already exist")

pet_prediction_app.include_router(auth_router, prefix="/api/auth")
pet_prediction_app.include_router(router, prefix="/api/predict")