face-shape / app.py
Stefano5699's picture
Llama deleted
d0c9299 verified
from fastapi import FastAPI, UploadFile, File
from transformers import pipeline
from fastai.vision.all import *
from PIL import Image
import os
access_token = os.getenv("access_token")
# NOTE - we configure docs_url to serve the interactive Docs at the root path
# of the app. This way, we can use the docs as a landing page for the app on Spaces.
app = FastAPI(docs_url="/")
pipe = pipeline("text2text-generation", model="google/flan-t5-small")
categories = ('Heart', 'Oblong', 'Oval', 'Round', 'Square')
learn = load_learner('model.pkl')
@app.get("/generate")
def generate(text: str):
"""
Using the text2text-generation pipeline from `transformers`, generate text
from the given input text. The model used is `google/flan-t5-small`, which
can be found [here](https://huggingface.co/google/flan-t5-small).
"""
output = pipe(text)
return {"output": output[0]["generated_text"]}
@app.post("/face-analyse")
async def face_analyse(file: UploadFile = File(...)):
# Read the uploaded file content
request_object_content = await file.read()
try:
# Attempt to open the image
img = Image.open(io.BytesIO(request_object_content))
except Exception as e:
return {"error": "Failed to open the image file. Make sure it is a valid image file."}
# Check if img is None or not
if img is None:
return {"error": "Failed to open the image file."}
try:
# Resize the image to 300x300 pixels
img = img.resize((300, 300))
except Exception as e:
return {"error": "Failed to resize the image."}
try:
# Assuming 'learn' is your image classifier model
pred, idx, probs = learn.predict(img)
except Exception as e:
return {"error": "Failed to make predictions."}
# Assuming categories is a list of category labels
return dict(zip(categories, map(float, probs)))