Spaces:
Sleeping
Sleeping
File size: 2,127 Bytes
7291333 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import tensorflow as tf
from PIL import Image
import numpy as np
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import io
import base64
import uvicorn
from fastapi.middleware.cors import CORSMiddleware
import pydantic
app = FastAPI()
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"]
)
app.mount("/static", StaticFiles(directory="static"), name="static")
@app.get("/")
def read_root():
return FileResponse("static/index.html")
# Load the model
model = tf.keras.models.load_model('model/k49_model.h5')
class ImageRequest(pydantic.BaseModel):
image: str
label: int
@app.post("/api/predict")
def predict(request: dict):
image = request["image"]
label = request["label"]
# Decode base64-encoded image
image_bytes = base64.b64decode(image["image"].split(",")[1])
# Convert image to preprocessed numpy array and save image to server
X = preprocess(image_bytes)
image = Image.fromarray(X[0].reshape(28, 28) * 255).convert("RGB")
image.save("image.png")
# Make prediction using the model
prediction = model.predict(X)
# Convert prediction to integer
prediction = int(np.argmax(prediction))
print("label: ", label)
print("prediction: ", prediction)
return {'prediction': prediction == int(label)}
def preprocess(image_bytes):
# Decode base64-encoded image and convert to PIL Image object
img = Image.open(io.BytesIO(image_bytes))
# Convert transparent background to white
if img.mode == 'RGBA':
img.load()
background = Image.new('RGB', img.size, (255, 255, 255))
background.paste(img, mask=img.split()[3])
img = background
# Convert to grayscale
img = img.convert('L')
# Invert colors to black background with white numbers
img = np.invert(np.array(img))
# Resize image to 28x28 pixels
img = Image.fromarray(img).resize((28, 28))
# Convert image to numpy array and normalize pixel values
X = np.array(img).reshape(1, 28, 28, 1) / 255
return X
if __name__ == '__main__':
uvicorn.run('app:app', host='localhost', port=8000, reload=True)
|