Spaces:
Sleeping
Sleeping
File size: 3,799 Bytes
38a3c61 1ef2263 38a3c61 4a30677 f73ccd7 1ef2263 f73ccd7 7e31555 1ef2263 f73ccd7 4a30677 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 120 121 122 123 124 |
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):
logging.info(f"Received request: {request.method} {request.url}")
api_key = request.headers.get("x-api-key")
if api_key is None or api_key not in VALID_API_KEYS:
logging.warning("Unauthorized access attempt.")
raise HTTPException(status_code=403, detail="Unauthorized")
response = await call_next(request)
return response
print(hf_token)
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)
# Return the probabilities as JSON
return JSONResponse(content={"confidence": confidence})
|