File size: 1,674 Bytes
4bd5ad3 6423eee 4bd5ad3 6423eee 4bd5ad3 6423eee 4bd5ad3 6423eee 4bd5ad3 | 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 | import json
from transformers import pipeline
import torch
class EndpointHandler:
"""
Handler customizado para o endpoint de inferência.
"""
def __init__(self, path=""):
# O 'path' é o diretório raiz. Construímos o caminho completo para a pasta do modelo final.
model_path = f"{path}/meu-modelo-ham10000-final"
# Determina se usará GPU ou CPU
device = 0 if torch.cuda.is_available() else -1
self.pipeline = pipeline(
"image-classification",
model=model_path, # <--- AGORA APONTAMOS PARA A PASTA EXATA!
device=device
)
print(f"Pipeline de classificação de imagem carregada com sucesso a partir de: {model_path}")
def __call__(self, data: dict) -> list:
"""
Esta função é chamada a cada requisição para a API.
'data' é o payload da requisição.
"""
# Extrai os inputs do payload da requisição
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", None)
print(f"Recebida requisição para analisar: {inputs}")
# Roda a inferência
try:
if parameters:
result = self.pipeline(inputs, **parameters)
else:
result = self.pipeline(inputs)
print(f"Resultado da análise: {result}")
return result
except Exception as e:
# Retorna um erro formatado se algo der errado durante a análise
error_message = f"Erro ao processar a requisição: {str(e)}"
print(error_message)
return [{"error": error_message}] |