Upload 2 files

main.py

@@ -0,0 +1,20 @@
+from fastapi import FastAPI, File
+from io import BytesIO
+from PIL import Image
+from predict import read_image, transformacao
+
+app = FastAPI()
+
+@app.get("/")
+async def root():
+    return {"message": "Idiot, you are in the wrong place!"}
+
+@app.post("/uploadfile/")
+async def create_upload_file(file: bytes = File(...)):
+
+    # read image
+    imagem = read_image(file)
+    # transform and prediction 
+    prediction = transformacao(imagem)
+
+    return prediction

predict.py

@@ -0,0 +1,36 @@
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from PIL import Image
+import torch
+import numpy as np
+from io import BytesIO  # Add this import statement
+
+processor = AutoImageProcessor.from_pretrained("dima806/medicinal_plants_image_detection")
+model = AutoModelForImageClassification.from_pretrained("dima806/medicinal_plants_image_detection")
+
+def read_image(file) -> Image.Image:
+    pil_image = Image.open(BytesIO(file))
+    return pil_image
+
+def transformacao(file: Image.Image):
+    inputs = processor(images=file, return_tensors="pt", padding=True)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits = outputs.logits
+    probabilities = logits.softmax(dim=1).squeeze()
+
+    # Get top 3 predictions
+    top3_probabilities, top3_indices = torch.topk(probabilities, 3)
+
+    labels = model.config.id2label
+
+    response = []
+    for prob, idx in zip(top3_probabilities, top3_indices):
+        resp = {}
+        resp["class"] = labels[idx.item()]
+        resp["confidence"] = f"{prob.item()*100:0.2f} %"
+        response.append(resp)
+
+    return response
+