Refactor API architecture with modular design and database integration

api/prediction.py

@@ -0,0 +1,65 @@
+import torch
+from fastapi import APIRouter, HTTPException
+from fastapi.responses import JSONResponse
+from PIL import Image
+from io import BytesIO
+import logging
+import httpx
+import asyncio
+from torchvision import transforms
+
+from schemas.requests import BatchPredictRequest
+from models.loader import get_model
+from steps.preprocess import process_image
+from config.settings import IMAGE_SIZE, NUM_THREADS
+
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+@router.post("/batch_predict")
+async def batch_predict(request: BatchPredictRequest):
+    """Endpoint pour prédire à partir de plusieurs images."""
+    model_name = request.modelName
+
+    try:
+        # Récupérer le modèle
+        model = get_model(model_name)
+        
+        semaphore = asyncio.Semaphore(NUM_THREADS)  # Limiter à 8 tâches simultanées
+        
+        async def process_single_image(image_url):
+            async with semaphore:
+                try:
+                    async with httpx.AsyncClient() as client:
+                        response = await client.get(image_url)
+                        image = Image.open(BytesIO(response.content))
+                except Exception:
+                    logger.error(f"Error downloading image from {image_url}")
+                    return {"imageUrl": image_url, "error": "Invalid image URL"}
+
+                # Prétraiter l'image
+                processed_image = process_image(image, size=IMAGE_SIZE)
+
+                # Convertir en tenseur
+                image_tensor = transforms.ToTensor()(processed_image).unsqueeze(0)
+
+                # Inférence
+                with torch.no_grad():
+                    outputs = model(image_tensor)
+                    probabilities = torch.nn.functional.softmax(outputs, dim=1)
+                    predicted_probabilities = probabilities.numpy().tolist()
+                    confidence = round(predicted_probabilities[0][1], 2)
+
+                return {"imageUrl": image_url, "confidence": confidence}
+
+        # Lancer les tâches en parallèle
+        tasks = [process_single_image(url) for url in request.imageUrls]
+        results = await asyncio.gather(*tasks)
+
+        return JSONResponse(content={"results": results})
+    
+    except KeyError:
+        raise HTTPException(status_code=404, detail="Model not found")
+    except Exception as e:
+        logger.error(f"Batch prediction error: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Batch prediction failed")

api/router.py

@@ -0,0 +1,23 @@
+from fastapi import APIRouter, Request, HTTPException
+import logging
+
+from api import prediction
+from config.settings import API_KEY
+
+logger = logging.getLogger(__name__)
+
+# Router principal
+router = APIRouter()
+
+# Middleware d'authentification
+async def verify_api_key(request: Request, call_next):
+    """Middleware pour vérifier la clé API dans les en-têtes."""
+    api_key = request.headers.get("x-api-key")
+    if api_key is None or api_key not in API_KEY.split(','):
+        logger.warning(f"Unauthorized API access attempt from {request.client.host}")
+        raise HTTPException(status_code=403, detail="Unauthorized")
+    response = await call_next(request)
+    return response
+
+# Inclure les routes des autres modules
+router.include_router(prediction.router, tags=["Prediction"])

config/settings.py

@@ -0,0 +1,31 @@
+import os
+import logging
+import torch
+
+# Configuration de base des logs
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Variables d'environnement obligatoires
+def get_env_or_fail(var_name: str) -> str:
+    """Récupérer une variable d'environnement ou échouer si elle n'est pas définie."""
+    value = os.environ.get(var_name)
+    if not value:
+        error_msg = f"{var_name} environment variable is not set or empty. API cannot start."
+        logger.error(error_msg)
+        raise RuntimeError(error_msg)
+    return value
+
+# Configuration API
+API_KEY = get_env_or_fail("api_key")
+HF_TOKEN = get_env_or_fail("api_read")
+RESOURCE_GROUP = get_env_or_fail("RESOURCE_GROUP")
+DATABASE_URL = get_env_or_fail("DATABASE_URL")
+
+# Log des paramètres importants (sans détails sensibles)
+logger.info(f"RESOURCE_GROUP set to: {RESOURCE_GROUP}")
+
+# Autres constantes
+IMAGE_SIZE = 256
+DEVICE = torch.device("cpu")  # Changer pour "cuda" si GPU disponible
+NUM_THREADS = 8  # Nombre de threads pour PyTorch

db/models.py

@@ -0,0 +1,48 @@
+import asyncpg
+import logging
+from typing import List, Dict, Any
+
+from config.settings import DATABASE_URL
+
+logger = logging.getLogger(__name__)
+
+async def fetch_models_for_group(resource_group: str) -> List[Dict[str, Any]]:
+    """Récupérer les détails des modèles depuis la base de données pour un groupe de ressources spécifique.
+    
+    Args:
+        resource_group: Identifiant du groupe de ressources
+        
+    Returns:
+        Liste de dictionnaires contenant les informations des modèles
+    
+    Raises:
+        Exception: Si une erreur se produit lors de la connexion ou de la requête
+    """
+    conn = None
+    try:
+        conn = await asyncpg.connect(DATABASE_URL)
+        logger.info(f"Successfully connected to database")
+
+        # Récupérer les modèles spécifiques à ce groupe
+        query = """
+            SELECT 
+                model_id, 
+                name, 
+                display_name, 
+                hf_repo_id, 
+                hf_subfolder, 
+                hf_filename
+            FROM models
+            WHERE hf_resource_group = $1
+        """
+        rows = await conn.fetch(query, resource_group)
+        logger.info(f"Found {len(rows)} models for group '{resource_group}'")
+        
+        return [dict(row) for row in rows]
+    except Exception as e:
+        logger.error(f"Database error: {e}", exc_info=True)
+        raise
+    finally:
+        if conn and not conn.is_closed():
+            await conn.close()
+            logger.debug("Database connection closed")

main.py

@@ -1,223 +1,43 @@
-import os
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-from transformers import pipeline
-from torchvision import transforms
-from PIL import Image
-import requests
-from io import BytesIO
-from steps.preprocess import process_image
-from huggingface_hub import hf_hub_download
-from architecture.resnet import ResNet
-import torch
 import logging
-from typing import List
-import httpx
-import asyncio
+from fastapi import FastAPI
 
-app = FastAPI()
+from api.router import router, verify_api_key
+from db.models import fetch_models_for_group
+from models.loader import load_models
+from config.settings import RESOURCE_GROUP
 
-image_size = 256
-hf_token = os.environ.get("api_read")
-VALID_API_KEYS = os.environ.get("api_key")
-INSTANCE_GROUP = os.environ.get("INSTANCE_GROUP")
-if INSTANCE_GROUP:
-    logging.info(f"INSTANCE_GROUP={INSTANCE_GROUP}")
-else:
-    logging.warning("INSTANCE_GROUP not set; all models will be loaded")
+# Configuration de base des logs
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
 
+# Créer l'application FastAPI
+app = FastAPI(
+    title="Tamis AI Inference API",
+    description="API pour l'inférence des modèles de classification d'objets",
+    version="0.1.0",
+)
 
-@app.middleware("http")
-async def verify_api_key(request, call_next):
-    api_key = request.headers.get("x-api-key")
-    if api_key is None or api_key not in VALID_API_KEYS:
-        raise HTTPException(status_code=403, detail="Unauthorized")
-    response = await call_next(request)
-    return response
-
-
-models_locations = [
-    # {
-    #     "repo_id": "TamisAI/category-lamp",
-    #     "subfolder": "maison-jansen/palmtree-152-0005-32-256",
-    #     "filename": "palmtree-jansen.pth",
-    # },
-    {
-        "repo_id": "TamisAI/category-lamp",
-        "subfolder": "maison-charles/corail-152-0001-32-256-L1",
-        "filename": "maison-charles-corail-L1.pth",
-    },
-    {
-        "repo_id": "TamisAI/category-lamp",
-        "subfolder": "michel-armand/flamme-152-0001A-32-256-L1",
-        "filename": "flamme-L1.pth",
-    },
-]
-
-device = torch.device("cpu")
-
-
-# Modèle de données pour les requêtes
-class PredictRequest(BaseModel):
-    imageUrl: str
-    modelName: str
-
-
-torch.set_num_threads(8)
-
-# Dictionnaire pour stocker les pipelines de modèles
-model_pipelines = {}
-
-# Create a single instance of the ResNet model
-base_model = ResNet("resnet152", num_output_neurons=2).to(device)
+# Ajouter middleware d'authentification
+app.middleware("http")(verify_api_key)
 
+# Inclure les routes
+app.include_router(router)
 
 @app.on_event("startup")
-async def load_models():
-    # Charger les modèles au démarrage
-    print(f"Loading models...{len(models_locations)}")
-
-    for model_location in models_locations:
-        try:
-            print(f"Loading model: {model_location['filename']}")
-            model_weight = hf_hub_download(
-                repo_id=model_location["repo_id"],
-                subfolder=model_location["subfolder"],
-                filename=model_location["filename"],
-                token=hf_token,
-            )
-            model = base_model.__class__("resnet152", num_output_neurons=2).to(device)
-            model.load_state_dict(
-                torch.load(model_weight, weights_only=True, map_location=device)
-            )
-            model.eval()
-            model_pipelines[model_location["filename"]] = model
-        except Exception as e:
-            print(f"Error loading model {model_location['filename']}: {e}")
-    print(f"Models loaded. {len(model_pipelines)}")
-
-
-@app.post("/predict")
-async def predict(request: PredictRequest):
-    image_url = request.imageUrl
-    model_name = request.modelName
-
-    # Télécharger l'image depuis l'URL
-    try:
-        response = requests.get(image_url)
-        image = Image.open(BytesIO(response.content))
-    except Exception as e:
-        raise HTTPException(status_code=400, detail="Invalid image URL")
-
-    # Vérifier si le modèle est chargé
-    if model_name not in model_pipelines:
-        raise HTTPException(status_code=404, detail="Model not found")
-
-    # Preprocess the image
-    processed_image = process_image(image, size=image_size)
-
-    # Convert to tensor
-    image_tensor = transforms.ToTensor()(processed_image).unsqueeze(0)
-
-    model = model_pipelines[model_name]
-
-    # Perform inference
-    with torch.no_grad():
-        outputs = model(image_tensor)
-        probabilities = torch.nn.functional.softmax(outputs, dim=1)
-        predicted_probabilities = probabilities.numpy().tolist()
-        confidence = round(predicted_probabilities[0][1], 2)
-    logging.info("confidence: %s", confidence)
-    # Return the probabilities as JSON
-    return JSONResponse(content={"confidence": confidence})
-
-
-class BatchPredictRequest(BaseModel):
-    imageUrls: List[str]
-    modelName: str
-
-
-# @app.post("/batch_predict")
-# async def batch_predict(request: BatchPredictRequest):
-#     model_name = request.modelName
-#     results = []
-
-#     # Verify if the model is loaded
-#     if model_name not in model_pipelines:
-#         raise HTTPException(status_code=404, detail="Model not found")
-
-#     model = model_pipelines[model_name]
-
-#     # Asynchronously process each image
-#     async with httpx.AsyncClient() as client:
-#         for image_url in request.imageUrls:
-#             try:
-#                 response = await client.get(image_url)
-#                 image = Image.open(BytesIO(response.content))
-#             except Exception as e:
-#                 results.append({"imageUrl": image_url, "error": "Invalid image URL"})
-#                 continue
-
-#             # Preprocess the image
-#             processed_image = process_image(image, size=image_size)
-
-#             # Convert to tensor
-#             image_tensor = transforms.ToTensor()(processed_image).unsqueeze(0)
-
-#             # Perform inference
-#             with torch.no_grad():
-#                 outputs = model(image_tensor)
-#                 probabilities = torch.nn.functional.softmax(outputs, dim=1)
-#                 predicted_probabilities = probabilities.numpy().tolist()
-#                 confidence = round(predicted_probabilities[0][1], 2)
-
-#             results.append({"imageUrl": image_url, "confidence": confidence})
-
-#     # Return the results as JSON
-#     return JSONResponse(content={"results": results})
-
-
-@app.post("/batch_predict")
-async def batch_predict(request: BatchPredictRequest):
-    model_name = request.modelName
-
-    # Verify if the model is loaded
-    if model_name not in model_pipelines:
-        raise HTTPException(status_code=404, detail="Model not found")
-
-    model = model_pipelines[model_name]
-    semaphore = asyncio.Semaphore(
-        8
-    )  # Limiter à 8 tâches simultanées pour éviter de surcharger la machine
-
-    async def process_single_image(image_url):
-        async with semaphore:
-            try:
-                async with httpx.AsyncClient() as client:
-                    response = await client.get(image_url)
-                    image = Image.open(BytesIO(response.content))
-            except Exception:
-                return {"imageUrl": image_url, "error": "Invalid image URL"}
-
-            # Preprocess the image
-            processed_image = process_image(image, size=image_size)
-
-            # Convert to tensor
-            image_tensor = transforms.ToTensor()(processed_image).unsqueeze(0)
-
-            # Perform inference
-            with torch.no_grad():
-                outputs = model(image_tensor)
-                probabilities = torch.nn.functional.softmax(outputs, dim=1)
-                predicted_probabilities = probabilities.numpy().tolist()
-                confidence = round(predicted_probabilities[0][1], 2)
-
-            return {"imageUrl": image_url, "confidence": confidence}
-
-    # Launch tasks in parallel
-    tasks = [process_single_image(url) for url in request.imageUrls]
-    results = await asyncio.gather(*tasks)
-
-    # Return the results as JSON
-    return JSONResponse(content={"results": results})
+async def startup():
+    """Initialiser l'API : charger les modèles depuis la base de données."""
+    logger.info("Starting API initialization...")
+    
+    # Charger les modèles depuis la base de données
+    models_data = await fetch_models_for_group(RESOURCE_GROUP)
+    await load_models(models_data)
+    
+    logger.info("API initialization complete.")
+
+@app.get("/health")
+async def health_check():
+    """Point d'entrée pour vérifier l'état de l'API."""
+    return {"status": "healthy"}

models/loader.py

@@ -0,0 +1,88 @@
+import torch
+import logging
+from typing import Dict, List, Any
+from huggingface_hub import hf_hub_download
+
+from config.settings import DEVICE, HF_TOKEN, NUM_THREADS
+from architecture.resnet import ResNet
+
+logger = logging.getLogger(__name__)
+
+# Configuration de PyTorch
+torch.set_num_threads(NUM_THREADS)
+
+# Instance de base pour le modèle ResNet
+base_model = ResNet("resnet152", num_output_neurons=2).to(DEVICE)
+
+# Dictionnaire global pour stocker les modèles chargés
+model_pipelines = {}
+
+async def load_models(models_data: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """Charger les modèles depuis Hugging Face à partir des données de la base de données.
+    
+    Args:
+        models_data: Liste de dictionnaires contenant les informations des modèles
+        
+    Returns:
+        Dictionnaire des modèles chargés
+        
+    Raises:
+        RuntimeError: Si aucun modèle n'est trouvé ou ne peut être chargé
+    """
+    logger.info(f"Attempting to load {len(models_data)} models...")
+    
+    if not models_data:
+        error_msg = "No models found. API cannot start without models."
+        logger.error(error_msg)
+        raise RuntimeError(error_msg)
+    
+    loaded_count = 0
+    for model_data in models_data:
+        try:
+            model_name = model_data['hf_filename']
+            logger.info(f"Loading model: {model_name} (repo: {model_data['hf_repo_id']}, subfolder: {model_data['hf_subfolder']})")
+            
+            model_weight = hf_hub_download(
+                repo_id=model_data['hf_repo_id'],
+                subfolder=model_data['hf_subfolder'],
+                filename=model_name,
+                token=HF_TOKEN,
+            )
+            
+            # Créer une nouvelle instance pour chaque modèle pour tenir ses poids spécifiques
+            model = base_model.__class__("resnet152", num_output_neurons=2).to(DEVICE)
+            model.load_state_dict(
+                torch.load(model_weight, weights_only=True, map_location=DEVICE)
+            )
+            model.eval()
+            model_pipelines[model_name] = model
+            loaded_count += 1
+        except Exception as e:
+            logger.error(f"Error loading model {model_data.get('hf_filename', 'N/A')}: {e}", exc_info=True)
+    
+    logger.info(f"Model loading finished. Successfully loaded {loaded_count}/{len(models_data)} models.")
+    
+    if loaded_count == 0:
+        error_msg = "Failed to load any models. API cannot start without models."
+        logger.error(error_msg)
+        raise RuntimeError(error_msg)
+    
+    return model_pipelines
+
+def get_model(model_name: str):
+    """Récupérer un modèle chargé par son nom.
+    
+    Args:
+        model_name: Nom du modèle à récupérer
+        
+    Returns:
+        Le modèle chargé
+        
+    Raises:
+        KeyError: Si le modèle n'est pas trouvé
+    """
+    if model_name not in model_pipelines:
+        logger.error(f"Model {model_name} not found in loaded models")
+        raise KeyError(f"Model {model_name} not found")
+    
+    return model_pipelines[model_name]

requirements.txt

@@ -7,4 +7,5 @@ torchvision
 huggingface_hub
 torch
 numpy
-httpx
+httpx
+asyncpg

schemas/requests.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+from typing import List
+
+class BatchPredictRequest(BaseModel):
+    """Modèle de requête pour des prédictions sur plusieurs images."""
+    imageUrls: List[str]
+    modelName: str

steps/preprocess.py

@@ -1,4 +1,3 @@
-import os
 from PIL import Image
 import numpy as np