Add potability IoT endpoint

README.md

@@ -1,4 +1,4 @@
-# WaterWise REST API ML Service
+# WaterWise ML Service
 
 This is the REST API service for the machine learning models for the WaterWise project.
 
@@ -7,6 +7,7 @@ This is the REST API service for the machine learning models for the WaterWise p
 - `POST /water-segmentation`: Segment/extract the water from the image. Returns an image.
 - `POST /clean-water`: Predict if the water is clean or not. Returns a float value from 0 to 1. Values above 0.5 indicate clean water.
 - `POST /clean-water/with-extraction`: Predict if the water is clean or not and extract the water from the image. Returns a float value from 0 to 1. Values above 0.5 indicate clean water.
+- `POST /potability-iot`: Predict if the water is drinkable or not (potability). Returns either 0 or 1. 1 indicate clean water.
 
 ## Usage
 
@@ -15,6 +16,7 @@ Prepare the models. Place the models in the [models](models) folder.
 ```bash
 wget "https://example.com/models/clean-water.h5" -O "clean-water.h5"
 wget "https://example.com/models/water-segmentation.pth" -O "water-segmentation.pth"
+wget "https://example.com/models/potability-iot.h5" -O "potability-iot.h5"
 ```
 
 Install the required packages using pip.

main.py

@@ -4,15 +4,15 @@
 import cv2
 from io import BytesIO
 
-from services import clean_water, water_segmentation
+from services import clean_water, water_segmentation, potability_iot
 
 
 app = flask.Flask(__name__)
 
 @app.route('/', methods=['GET'])
 def index():
     return flask.jsonify({
-        "message": "WaterWise REST API ML Service",
+        "message": "WaterWise ML Service",
         "data": {
             "time": datetime.datetime.now()
         }
@@ -87,6 +87,39 @@ def predict_clean_water_full():
         }
     })
 
+@app.route('/potability-iot', methods=['POST'])
+def predict_potability_iot():
+
+    if flask.request.form:
+        data = flask.request.form
+    else:
+        data = flask.request.json
+
+    print(data)
+    print(data['solids'])
+    print(data['turbidity'])
+    print(data['chloramines'])
+    print(data['organic_carbon'])
+    print(data['sulfate'])
+    print(data['ph'])
+
+
+    prediction = potability_iot.predict(
+        solids=float(data['solids']),
+        turbidity=float(data['turbidity']),
+        chloramines=float(data['chloramines']),
+        organic_carbon=float(data['organic_carbon']),
+        sulfate=float(data['sulfate']),
+        ph=float(data['ph'])
+    )
+
+    return flask.jsonify({
+        "message": "Prediction successful",
+        "data": {
+            "prediction": float(prediction[0][0])
+        }
+    })
+
 def create_app():
    return app
 

models/README.md

@@ -3,4 +3,5 @@
 Two files should reside in this folder:
 
 - `clean-water.h5`: The trained model for the clean water detection
-- `water-segmentation.pth`: The trained model for the water segmentation
+- `water-segmentation.pth`: The trained model for the water segmentation
+- `potability-iot.h5`: The trained model for the water potability prediction using IoT data

services/potability_iot.py

@@ -0,0 +1,11 @@
+import numpy as np
+import tensorflow as tf
+
+import cv2
+
+model = tf.keras.models.load_model("models/potability-iot.h5")
+
+def predict(solids: float, turbidity: float, chloramines: float, organic_carbon: float, sulfate: float, ph: float):
+    prediction = model.predict(np.array([[solids, turbidity, chloramines, organic_carbon, sulfate, ph]]))
+
+    return prediction