|
|
|
|
|
|
|
|
from fastapi import FastAPI, File, UploadFile |
|
|
import uvicorn |
|
|
import numpy as np |
|
|
from io import BytesIO |
|
|
from PIL import Image |
|
|
import tensorflow as tf |
|
|
|
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
|
|
|
class ImageClassifier: |
|
|
def __init__(self, model_path): |
|
|
self.MODEL = tf.keras.models.load_model(model_path) |
|
|
self.CLASS_NAMES = ["Not Coffee Land", "Coffee Land"] |
|
|
|
|
|
|
|
|
|
|
|
def read_file_as_image(self, data) -> np.ndarray: |
|
|
image = np.array(Image.open(BytesIO(data))) |
|
|
return image |
|
|
|
|
|
|
|
|
def predict(self, file: UploadFile): |
|
|
image = self.read_file_as_image(file.file.read()) |
|
|
img_batch = np.expand_dims(image, 0) |
|
|
|
|
|
predictions = self.MODEL.predict(img_batch) |
|
|
predicted_class = self.CLASS_NAMES[np.argmax(predictions[0])] |
|
|
|
|
|
return {'class': predicted_class} |
|
|
|
|
|
classifier = ImageClassifier("model.h5") |
|
|
|
|
|
|
|
|
@app.get("/ping") |
|
|
async def ping(): |
|
|
return "Hello World" |
|
|
|
|
|
@app.post("/predict") |
|
|
async def predict(file: UploadFile = File(...)): |
|
|
result = classifier.predict(file) |
|
|
return result |
|
|
|
|
|
if __name__ == "__main__": |
|
|
uvicorn.run(app, host='localhost', port=8080) |
|
|
|