Added ai drink detection

backend/.dockerignore

@@ -0,0 +1 @@
+ml/

backend/.gitignore

@@ -156,4 +156,5 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-*.db
+*.db
+ml/

backend/TaskScheduler.py

@@ -23,6 +23,12 @@ def add_Daily_Task(self, task, *args) -> None:
         else:
             schedule.every().day.at("00:01").do(task)
 
+    def add_Weekly_Task(self, task, *args) -> None:
+        if len(args) > 0:
+            schedule.every(7).days.do(task).do(task, args)
+        else:
+            schedule.every(7).days.do(task)
+
     def add_5min_Task(self, task, *args) -> None:
         if len(args) > 0:
             schedule.every(5).minutes.do(task, args)

backend/ai.py

@@ -0,0 +1,167 @@
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers, regularizers
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+
+
+# Set parameters
+image_size = (224, 224)
+dataset_path = "./ml/drinks-combined"
+model_path = "./ml/drinks_bottle_classifier-m6.keras"
+batch_size = 32
+
+# Training generator with augmentation (no rescaling)
+train_datagen = ImageDataGenerator(
+    validation_split=0.2,
+    rotation_range=20,
+    width_shift_range=0.2,
+    height_shift_range=0.2,
+    shear_range=0.2,
+    zoom_range=0.2,
+    horizontal_flip=True,
+    brightness_range=[0.8, 1.2]
+)
+
+# Validation generator without augmentation (no rescaling)
+val_datagen = ImageDataGenerator(
+    validation_split=0.2
+)
+
+train_generator = train_datagen.flow_from_directory(
+    dataset_path,
+    target_size=image_size,
+    batch_size=batch_size,
+    class_mode='categorical',
+    subset='training'
+)
+
+val_generator = val_datagen.flow_from_directory(
+    dataset_path,
+    target_size=image_size,
+    batch_size=batch_size,
+    class_mode='categorical',
+    subset='validation'
+)
+
+# Custom callback to plot training and validation accuracy after each epoch
+
+
+class AccuracyPlotCallback(keras.callbacks.Callback):
+    def __init__(self, output_file="accuracy_plot.png"):
+        super().__init__()
+        self.history_acc = []
+        self.history_val_acc = []
+        self.output_file = output_file
+
+    def on_epoch_end(self, epoch, logs=None):
+        self.history_acc.append(logs.get('accuracy'))
+        self.history_val_acc.append(logs.get('val_accuracy'))
+        print(
+            f"Epoch {epoch+1} - Accuracy: {self.history_acc[-1]:.4f}, Val Accuracy: {self.history_val_acc[-1]:.4f}")
+        plt.figure()
+        plt.plot(self.history_acc, label='accuracy', color='blue')
+        plt.plot(self.history_val_acc, label='val_accuracy', color='orange')
+        plt.xlabel('Epoch')
+        plt.ylabel('Accuracy')
+        plt.legend()
+        plt.savefig(self.output_file)
+        plt.close()
+
+    def on_train_end(self, logs=None):
+        plt.figure()
+        plt.plot(self.history_acc, label='accuracy', color='blue')
+        plt.plot(self.history_val_acc, label='val_accuracy', color='orange')
+        plt.xlabel('Epoch')
+        plt.ylabel('Accuracy')
+        plt.legend()
+        plt.savefig(self.output_file)
+        plt.close()
+
+# Define a model using pre-trained MobileNetV2
+
+
+def create_model():
+    # Load the MobileNetV2 model with pre-trained ImageNet weights
+    base_model = keras.applications.MobileNetV2(
+        input_shape=(image_size[0], image_size[1], 3),
+        include_top=False,
+        weights='imagenet'
+    )
+    base_model.trainable = False  # Freeze the pre-trained layers
+
+    inputs = keras.Input(shape=(image_size[0], image_size[1], 3))
+    # Apply MobileNetV2 preprocessing to convert [0,255] to [-1,1]
+    x = keras.applications.mobilenet_v2.preprocess_input(inputs)
+    x = base_model(x, training=False)
+    x = keras.layers.GlobalAveragePooling2D()(x)
+    x = keras.layers.Dropout(0.5)(x)
+    outputs = keras.layers.Dense(
+        len(train_generator.class_indices), activation='softmax')(x)
+    model = keras.Model(inputs, outputs)
+    return model
+
+
+def learn(epochs=30,path=None):
+    model = create_model()
+    model.compile(optimizer='adam',
+                  loss='categorical_crossentropy',
+                  metrics=['accuracy'])
+
+    lr_scheduler = keras.callbacks.ReduceLROnPlateau(
+        monitor='val_loss', patience=3, factor=0.5)
+    early_stopping = keras.callbacks.EarlyStopping(
+        monitor='val_loss', patience=5, restore_best_weights=True)
+    acc_plot_callback = AccuracyPlotCallback(
+        output_file="./plots/accuracy_plot.png")
+    os.makedirs("./plots", exist_ok=True)
+
+    history = model.fit(
+        train_generator,
+        validation_data=val_generator,
+        epochs=epochs,
+        callbacks=[lr_scheduler, early_stopping, acc_plot_callback]
+    )
+    model.save(model_path if path is None else path)
+    return model
+
+
+def predict_image(image_path, model):
+    from tensorflow.keras.preprocessing import image
+    img = image.load_img(image_path, target_size=image_size)
+    img_array = image.img_to_array(img)
+    # No rescaling here: MobileNetV2 preprocesses input in the model
+    img_array = np.expand_dims(img_array, axis=0)
+    predictions = model.predict(img_array)[0]
+
+    top_4_indices = np.argsort(predictions)[-4:][::-1]
+
+    top_4_labels = [list(train_generator.class_indices.keys())[i]
+                    for i in top_4_indices]
+    top_4_probs = [predictions[i] for i in top_4_indices]
+    return list(zip(top_4_labels, top_4_probs))
+
+
+model = None
+# model = learn(30)
+
+
+def learn_and_set(epochs=30):
+    global model
+    new_model = learn(epochs=epochs)
+    model = new_model
+
+learn_ai = os.environ.get("AI_LEARN") == "true" if os.environ.get(
+    "AI_LEARN") else False
+
+if learn_ai:
+    learn(path=input("Model output path:")+"/drinks_bottle_classifier-m6.keras")
+
+if model is None:
+    model = keras.models.load_model(model_path)
+
+
+def predict_jpg(filename):
+    return predict_image(filename, model)

backend/database/Queries.py

@@ -770,6 +770,22 @@ def enable_disable_pretix_user(self):
                 # Add user to database
                 self.add_user(name, 0, util.standard_user_password,
                               name, hidden=not checked_in)
+
+    def get_drink_id_by_closest_name(self, name):
+        drinks = self.session.query(Drink).all()
+        best_match = None
+        best_ratio = 0
+        for drink in drinks:
+            ratio = SequenceMatcher(None, name, drink.name).ratio()
+            if ratio > best_ratio:
+                best_match = drink
+                best_ratio = ratio
+
+        return best_match.id
+
+    def get_drink_from_id(self, drink_id):
+        drink: Drink = self.session.query(Drink).filter_by(id=drink_id).first()
+        return drink
 
     def is_admin(self, member_id):
         member: Member = self.session.query(

backend/main.py

@@ -19,6 +19,11 @@
 import flask
 import mail
 import secrets
+import base64
+import secrets
+
+if util.ai_enabled:
+    import ai
 
 api_bp = flask.Blueprint("api", __name__, url_prefix="/api/")
 api = Api(api_bp, doc='/docu/', base_url='/api')
@@ -38,6 +43,8 @@
     if util.pretix_url is not None:
         db.enable_disable_pretix_user()
         taskScheduler.add_5min_Task(db.enable_disable_pretix_user)
+    # if util.ai_enabled:
+    #    taskScheduler.add_Weekly_Task(ai.learn_and_set)
     taskScheduler.start()
 
 
@@ -452,6 +459,110 @@ def get(self):
         return util.build_response(db.get_drink_categories())
 
 
+@api.route('/drinks/ai/train')
+class drink_training(Resource):
+    @authenticated
+    def post(self):
+        """
+        Receive base64 encoded image and store the image on the server
+        """
+        image = request.json["image"]
+        image = image.split(",")[1]
+        image = image.encode()
+
+        # If directory does not exist, create it
+        if not os.path.exists("ml/training"):
+            os.makedirs("ml/training")
+
+        # current time with milliseconds as filename
+        time = datetime.now().strftime("%Y-%m-%d-%H-%M-%S-%f")
+
+        with open("ml/training/"+time+".jpg", "wb") as fh:
+            fh.write(base64.decodebytes(image))
+            return util.build_response("Image received")
+
+        return util.build_response(code=500)
+
+
+@api.route('/drinks/ai/recognition')
+class drink_training(Resource):
+    @authenticated
+    def post(self):
+        """
+        Receive base64 encoded image and store the image on the server
+        """
+        image = request.json["image"]
+        image = image.split(",")[1]
+        image = image.encode()
+
+        # If directory does not exist, create it
+        if not os.path.exists("ml/recognition"):
+            os.makedirs("ml/recognition")
+
+        image_path = secrets.token_urlsafe(24) + ".jpg"
+
+        with open("ml/recognition/"+image_path, "wb") as fh:
+            fh.write(base64.decodebytes(image))
+
+        predicted_labels = ai.predict_jpg("ml/recognition/"+image_path)
+
+        # Delete the image after prediction
+        os.remove("ml/recognition/"+image_path)
+
+        output = []
+
+        print(predicted_labels)
+
+        if predicted_labels[0][0] == "nothing":
+            return util.build_response([], code=404)
+
+        if predicted_labels[0][1] > 0.9:
+            # Search for all labels the best matching drink
+            for label in predicted_labels:
+                drink_id_to_add = db.get_drink_id_by_closest_name(label[0])
+                if drink_id_to_add is not None:
+                    output.append(drink_id_to_add)
+
+            return util.build_response(output)
+
+        return util.build_response([], code=404)
+
+
+@api.route('/drinks/ai/training/user')
+class user_identified_data(Resource):
+
+    @authenticated
+    def post(self):
+        """
+        Receive base64 encoded image and store the image on the server
+        """
+        image = request.json["image"]
+        image = image.split(",")[1]
+        image = image.encode()
+
+        # If directory does not exist, create it
+        if not os.path.exists("ml/training-user"):
+            os.makedirs("ml/training-user")
+
+        # current time with milliseconds as filename
+        time = datetime.now().strftime("%Y-%m-%d-%H-%M-%S-%f")
+
+        drink_id = request.json["drinkID"]
+        drink: Drink = db.get_drink_from_id(drink_id)
+
+        drink_directory = "ml/training-user/"+drink.name
+
+        # check if drink directory exists
+        if not os.path.exists(drink_directory):
+            os.makedirs(drink_directory)
+
+        with open(drink_directory+"/"+time+".jpg", "wb") as fh:
+            fh.write(base64.decodebytes(image))
+            return util.build_response("Image received")
+
+        return util.build_response(code=500)
+
+
 @api.route('/drinks/<int:drink_id>/price')
 class set_drink_price(Resource):
     @admin
@@ -460,6 +571,9 @@ def post(self, drink_id):
         """
         Set the price of a drink
         """
+        if request.json["amount"] is None:
+            return util.build_response("Price cannot be empty", code=406)
+
         db.change_drink_price(drink_id, request.json["amount"])
         return util.build_response("Price changed")
 

backend/requirements.txt

@@ -8,3 +8,9 @@ psycopg2
 waitress
 requests
 werkzeug
+tensorflow
+numpy
+matplotlib
+pillow
+scipy
+autokeras

backend/util.py

@@ -4,6 +4,7 @@
 import datetime
 import time
 import requests
+import random
 
 cookie_expire = int(os.environ.get("COOKIE_EXPIRE_TIME")) * \
     60*60 if os.environ.get("COOKIE_EXPIRE_TIME") else 60**3
@@ -83,6 +84,10 @@
 pretix_api_token = os.environ.get("PRETIX_API_TOKEN") if os.environ.get(
     "PRETIX_API_TOKEN") else None
 
+ai_enabled = os.environ.get("AI_ENABLED") == "true" if os.environ.get(
+    "AI_ENABLED") else False
+ai_model_path = os.environ.get("AI_MODEL_PATH") if os.environ.get(
+    "AI_MODEL_PATH") else "drinks_bottle_classifier-m6.keras"
 
 tempfile_path = "tempfiles"
 backup_file_name = "backup.json"