salek877 commited on
Commit
606917a
1 Parent(s): 3487197
Files changed (1) hide show
  1. SATIN API/main.py +48 -0
SATIN API/main.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # main.py
2
+
3
+ from fastapi import FastAPI, File, UploadFile
4
+ import uvicorn
5
+ import numpy as np
6
+ from io import BytesIO
7
+ from PIL import Image
8
+ import tensorflow as tf
9
+
10
+
11
+ app = FastAPI()
12
+
13
+
14
+ class ImageClassifier:
15
+ def __init__(self, model_path):
16
+ self.MODEL = tf.keras.models.load_model(model_path) # Load model
17
+ self.CLASS_NAMES = ["Not Coffee Land", "Coffee Land"]
18
+
19
+
20
+ # Process input data
21
+ def read_file_as_image(self, data) -> np.ndarray:
22
+ image = np.array(Image.open(BytesIO(data)))
23
+ return image
24
+
25
+ # Return output
26
+ def predict(self, file: UploadFile):
27
+ image = self.read_file_as_image(file.file.read())
28
+ img_batch = np.expand_dims(image, 0)
29
+
30
+ predictions = self.MODEL.predict(img_batch)
31
+ predicted_class = self.CLASS_NAMES[np.argmax(predictions[0])]
32
+
33
+ return {'class': predicted_class}
34
+
35
+ classifier = ImageClassifier("model.h5")
36
+
37
+ # Created this for testing
38
+ @app.get("/ping")
39
+ async def ping():
40
+ return "Hello World"
41
+
42
+ @app.post("/predict")
43
+ async def predict(file: UploadFile = File(...)):
44
+ result = classifier.predict(file)
45
+ return result
46
+
47
+ if __name__ == "__main__":
48
+ uvicorn.run(app, host='localhost', port=8080)