File size: 1,545 Bytes
04b0935
 
 
 
 
 
 
 
8491fd6
04b0935
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b98102b
8491fd6
 
 
 
 
 
 
 
 
04b0935
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
48
49
50
51
52
53
54
55
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from tensorflow import keras
import tensorflow as tf
import os
import numpy as np
from PIL import Image
from dotenv import load_dotenv
from huggingface_hub import hf_hub_download

load_dotenv()

app = FastAPI()

# Make sure it allows all the requests for the CORSMiddleware.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

ANIMALS = ["Cat", "Dog", "Panda"]

os.makedirs("./model", exist_ok=True)

# Download your SavedModel from the Hub
repo_id = "drgou/howest-deployathome"
hf_hub_download(repo_id, filename="config.json", repo_type="model", local_dir="./model")
hf_hub_download(repo_id, filename="metadata.json", repo_type="model", local_dir="./model")
hf_hub_download(repo_id, filename="model.weights.h5", repo_type="model", local_dir="./model")

# Load it
model = tf.keras.models.load_model("./model")


@app.post('/upload/image')
async def uploadImage(img: UploadFile = File(...)):
    original_image = Image.open(img.file)
    resized_image = original_image.resize((64,64))
    images_to_predict = np.expand_dims(np.array(resized_image), axis=0)
    predictions = model.predict(images_to_predict)
    prediction_probabilities = predictions
    classifications = prediction_probabilities.argmax(axis=1)

    print(predictions)
    print(classifications)
    print(classifications.tolist()[0])

    return ANIMALS[classifications.tolist()[0]]