Update app.py

app.py

@@ -8,6 +8,7 @@ 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):
@@ -21,6 +22,7 @@ def generate(text: str):
 
 @app.post("/uploadfile/")
 async def create_upload_file(file: UploadFile):
-    return {"test":"test123"}
+    pred, idx, probs = learn.predict(img)
+    return dict(zip(categories, map(float, probs)))
 #    pred, idx, probs = learn.predict(img)
 #    return dict(zip(categories, map(float, probs)))