Update app.py

app.py

@@ -1,114 +1,122 @@
+# ================================
+# ENVIRONMENT FIXES (MUST BE FIRST)
+# ================================
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
 os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
+# ================================
+# IMPORTS
+# ================================
 import gradio as gr
 import tensorflow as tf
 import numpy as np
 import pickle
 from PIL import Image
 
-# ===============================
-# CONSTANTS
-# ===============================
+# ================================
+# CONSTANTS (MUST MATCH TRAINING)
+# ================================
 IMG_SIZE = (128, 128)
 MAX_LEN = 50
-NUM_CLASSES = 4
+MAX_WORDS = 20000
 LABELS = ["Critical", "High", "Medium", "Low"]
 
-# ===============================
-# IMAGE ENCODER (MATCH TRAINING)
-# ===============================
-def build_image_encoder():
-    inp = tf.keras.Input(shape=(128,128,3), name="image_input")
-
-    x = tf.keras.layers.Conv2D(32, 3, activation="relu")(inp)
-    x = tf.keras.layers.MaxPooling2D()(x)
-    x = tf.keras.layers.Conv2D(64, 3, activation="relu")(x)
-    x = tf.keras.layers.MaxPooling2D()(x)
-    x = tf.keras.layers.Conv2D(128, 3, activation="relu")(x)
-    x = tf.keras.layers.MaxPooling2D()(x)
-    x = tf.keras.layers.Flatten()(x)
-    x = tf.keras.layers.Dense(128, activation="relu", name="image_embedding")(x)
-
-    return tf.keras.Model(inp, x)
-
-# ===============================
-# TEXT ENCODER (MATCH TRAINING)
-# ===============================
-def build_text_encoder(vocab_size=20000):
-    inp = tf.keras.Input(shape=(MAX_LEN,), name="text_input")
-
-    x = tf.keras.layers.Embedding(vocab_size, 128)(inp)
-    x = tf.keras.layers.LSTM(64)(x)
-    x = tf.keras.layers.Dense(64, activation="relu", name="text_embedding")(x)
-
-    return tf.keras.Model(inp, x)
-
-# ===============================
-# FUSION MODEL
-# ===============================
-image_encoder = build_image_encoder()
-text_encoder = build_text_encoder()
-
-combined = tf.keras.layers.Concatenate()([
-    image_encoder.output,
-    text_encoder.output
-])
-
-x = tf.keras.layers.Dense(128, activation="relu")(combined)
-x = tf.keras.layers.Dense(64, activation="relu")(x)
-output = tf.keras.layers.Dense(NUM_CLASSES, activation="softmax")(x)
+# ================================
+# LOAD TOKENIZER
+# ================================
+with open("tokenizer.pkl", "rb") as f:
+    tokenizer = pickle.load(f)
+
+print("✅ Tokenizer loaded")
+
+# ================================
+# BUILD FUSION MODEL ARCHITECTURE
+# ================================
+
+# ---- Image encoder ----
+image_input = tf.keras.Input(shape=(128, 128, 3), name="image_input")
+
+x = tf.keras.layers.Conv2D(32, (3, 3), activation="relu")(image_input)
+x = tf.keras.layers.MaxPooling2D((2, 2))(x)
+x = tf.keras.layers.Conv2D(64, (3, 3), activation="relu")(x)
+x = tf.keras.layers.MaxPooling2D((2, 2))(x)
+x = tf.keras.layers.Conv2D(128, (3, 3), activation="relu")(x)
+x = tf.keras.layers.MaxPooling2D((2, 2))(x)
+x = tf.keras.layers.Flatten()(x)
+image_features = tf.keras.layers.Dense(128, activation="relu")(x)
+
+# ---- Text encoder ----
+text_input = tf.keras.Input(shape=(MAX_LEN,), name="text_input")
+
+y = tf.keras.layers.Embedding(MAX_WORDS, 128)(text_input)
+y = tf.keras.layers.LSTM(64)(y)
+text_features = tf.keras.layers.Dense(64, activation="relu")(y)
+
+# ---- Fusion head ----
+combined = tf.keras.layers.Concatenate()([image_features, text_features])
+z = tf.keras.layers.Dense(128, activation="relu")(combined)
+z = tf.keras.layers.Dense(64, activation="relu")(z)
+output = tf.keras.layers.Dense(4, activation="softmax")(z)
 
 fusion_model = tf.keras.Model(
-    inputs=[image_encoder.input, text_encoder.input],
+    inputs=[image_input, text_input],
     outputs=output
 )
 
-# ===============================
-# LOAD WEIGHTS (SAFE)
-# ===============================
+# ================================
+# LOAD WEIGHTS (THIS IS THE KEY)
+# ================================
 fusion_model.load_weights("fusion_weights.weights.h5")
 print("✅ Fusion weights loaded")
 
-# ===============================
-# LOAD TOKENIZER
-# ===============================
-with open("tokenizer.pkl", "rb") as f:
-    tokenizer = pickle.load(f)
-
-# ===============================
-# PREPROCESS
-# ===============================
-def preprocess_image(image):
-    image = image.convert("RGB").resize(IMG_SIZE)
-    img = np.array(image) / 255.0
-    return np.expand_dims(img, 0)
-
-def preprocess_text(text):
-    seq = tokenizer.texts_to_sequences([text or ""])
-    return tf.keras.preprocessing.sequence.pad_sequences(seq, maxlen=MAX_LEN)
-
-# ===============================
-# PREDICT
-# ===============================
+# ================================
+# PREPROCESSING
+# ================================
+def preprocess_image(image: Image.Image):
+    image = image.convert("RGB")
+    image = image.resize(IMG_SIZE)
+    img = np.array(image, dtype=np.float32) / 255.0
+    return np.expand_dims(img, axis=0)
+
+def preprocess_text(text: str):
+    if text is None:
+        text = ""
+    seq = tokenizer.texts_to_sequences([text])
+    return tf.keras.preprocessing.sequence.pad_sequences(
+        seq, maxlen=MAX_LEN
+    )
+
+# ================================
+# PREDICTION FUNCTION
+# ================================
 def predict_ticket(image, text):
     if image is None:
-        return {}
+        return {"Critical": 0, "High": 0, "Medium": 0, "Low": 0}
 
     img = preprocess_image(image)
     txt = preprocess_text(text)
 
     probs = fusion_model.predict([img, txt], verbose=0)[0]
-    return dict(zip(LABELS, map(float, probs)))
 
-# ===============================
-# UI
-# ===============================
+    return {
+        "Critical": float(probs[0]),
+        "High": float(probs[1]),
+        "Medium": float(probs[2]),
+        "Low": float(probs[3]),
+    }
+
+# ================================
+# GRADIO UI
+# ================================
 interface = gr.Interface(
     fn=predict_ticket,
-    inputs=[gr.Image(type="pil"), gr.Textbox(lines=3)],
-    outputs=gr.Label(num_top_classes=4),
-    title="Ticket Severity Classifier"
+    inputs=[
+        gr.Image(type="pil", label="📤 Upload Ticket Screenshot"),
+        gr.Textbox(lines=4, label="✍️ Ticket Description")
+    ],
+    outputs=gr.Label(num_top_classes=4, label="🚨 Predicted Severity"),
+    title="🎫 Ticket Severity Classification",
 )
 
-interface.launch()
+interface.launch(server_name="0.0.0.0", server_port=7860)