Spaces:
Sleeping
Sleeping
File size: 3,701 Bytes
38a3c61 1ef2263 38a3c61 4a30677 f73ccd7 c55bc6a f73ccd7 38a3c61 6d6e2b4 38a3c61 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
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
app = FastAPI()
image_size = 256
hf_token = os.environ.get("api_read")
VALID_API_KEYS = os.environ.get("api_key")
@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",
"filename": "maison-charles-corail.pth",
},
]
device = torch.device("cpu")
# Modèle de données pour les requêtes
class PredictRequest(BaseModel):
imageUrl: str
modelName: str
# 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)
@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})
|