AIIA_docker / app.py
mat27's picture
Update app.py
00991f2
raw
history blame
375 Bytes
from fastapi import FastAPI
from transformers import pipeline
import numpy as np
import os
os.environ['TRANSFORMERS_CACHE'] = '/blabla/cache/'
# Create a new FastAPI app instance
app = FastAPI()
pipe = pipeline(model="mat27/medmnistPrueba")
def generate(image: np.array):
"""
Input: imagen
"""
output = pipe(image)
return {"output": output[0]}