Update app.py

app.py

@@ -1,162 +1,162 @@
-import os
-import numpy as np
-import tensorflow as tf
-from flask import Flask, request, render_template, send_from_directory
-from werkzeug.utils import secure_filename
-from PIL import UnidentifiedImageError # To catch corrupted image errors
-
-# --- 1. CONFIGURATION ---
-app = Flask(__name__)
-UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-
-# Define paths for the .h5 models
-BREED_MODEL_PATH = 'dog_breed_project_model.h5'      # <-- UPDATED FILENAME
-GATEKEEPER_MODEL_PATH = 'gatekeeper_model.h5'        # <-- UPDATED FILENAME
-CLASS_NAMES_PATH = 'class_names.txt'
-IMG_SIZE = 224
-
-# --- 2. LOAD MODELS AND CLASS NAMES (ONCE AT STARTUP) ---
-print("--- Loading models and class names... ---")
-try:
-    # Load the expert model for breed classification
-    breed_model = tf.keras.models.load_model(BREED_MODEL_PATH)
-    print("✅ Dog Breed (Expert) model loaded.")
-
-    # Load the gatekeeper model for dog vs. not-dog classification
-    gatekeeper_model = tf.keras.models.load_model(GATEKEEPER_MODEL_PATH)
-    print("✅ Gatekeeper (Dog vs. Not-Dog) model loaded.")
-
-    # Load class names
-    with open(CLASS_NAMES_PATH, 'r') as f:
-        class_names = [line.strip() for line in f.readlines()]
-    print(f"✅ Class names loaded. Found {len(class_names)} classes.")
-
-except Exception as e:
-    print(f"❌ Error loading models: {e}")
-    # Define dummy models so the app doesn't crash on startup
-    breed_model = None
-    gatekeeper_model = None
-
-print("--- Setup complete ---")
-
-
-# --- 3. IMAGE PREPROCESSING FUNCTION (CONSISTENT WITH TRAINING) ---
-def preprocess_uploaded_image(filepath, img_size):
-    """
-    Loads, decodes, and preprocesses an uploaded image for both models.
-    This function handles different file types, grayscale images, and aspect ratios.
-    """
-    try:
-        # Read the file and decode it as a 3-channel (RGB) image
-        img = tf.io.read_file(filepath)
-        img = tf.image.decode_image(img, channels=3, expand_animations=False)
-
-        # Pad to a square aspect ratio without distortion
-        img = tf.image.resize_with_pad(img, img_size, img_size)
-        
-        # Expand dimensions to create a batch of 1
-        img_batch = tf.expand_dims(img, 0)
-
-        # Preprocess for each model's specific requirements
-        gatekeeper_input = tf.keras.applications.mobilenet_v2.preprocess_input(tf.identity(img_batch))
-        breed_model_input = tf.keras.applications.resnet_v2.preprocess_input(tf.identity(img_batch))
-        
-        return gatekeeper_input, breed_model_input
-
-    except (UnidentifiedImageError, tf.errors.InvalidArgumentError):
-        # Handle cases where the file is not a valid image
-        return None, None
-    except Exception as e:
-        print(f"An unexpected error occurred during preprocessing: {e}")
-        return None, None
-
-
-# --- 4. FLASK ROUTES ---
-@app.route('/', methods=['GET', 'POST'])
-def index():
-    if request.method == 'POST':
-        if not all([breed_model, gatekeeper_model]):
-             return render_template('index.html', error="Models are not loaded. Please check the server logs.")
-
-        if 'file' not in request.files:
-            return render_template('index.html', error="No file part in the request.")
-        
-        file = request.files['file']
-        if file.filename == '':
-            return render_template('index.html', error="No file selected.")
-
-        if file:
-            filename = secure_filename(file.filename)
-            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-            file.save(filepath)
-
-            # --- PREDICTION PIPELINE ---
-            gatekeeper_img, breed_img = preprocess_uploaded_image(filepath, IMG_SIZE)
-
-            if gatekeeper_img is None:
-                return render_template('index.html', error="Invalid or corrupted image file. Please try another.")
-
-            # Step 1: Use the Gatekeeper to check if the image is a dog
-            gatekeeper_pred = gatekeeper_model.predict(gatekeeper_img)[0][0]
-            
-            if gatekeeper_pred > 0.5:
-                # -----------------------------------------------------------------
-                # START: UPGRADED LOGIC WITH TEST-TIME AUGMENTATION (TTA)
-                # -----------------------------------------------------------------
-                
-                # Create 4 augmented versions of the image (0, 90, 180, 270 degrees)
-                images_to_predict = [
-                    breed_img,
-                    tf.image.rot90(breed_img, k=1), # 90 degrees
-                    tf.image.rot90(breed_img, k=2), # 180 degrees (upside down)
-                    tf.image.rot90(breed_img, k=3)  # 270 degrees
-                ]
-                
-                # Stack the images into a single batch
-                tta_batch = tf.concat(images_to_predict, axis=0)
-                
-                # Get predictions for all 4 images in one go
-                batch_predictions = breed_model.predict(tta_batch)
-                
-                # Average the predictions to get the final, robust result
-                breed_predictions = tf.reduce_mean(batch_predictions, axis=0)
-                
-                # -----------------------------------------------------------------
-                # END: UPGRADED LOGIC WITH TEST-TIME AUGMENTATION (TTA)
-                # -----------------------------------------------------------------
-
-                # Get top 3 predictions from the averaged result
-                top_k_values, top_k_indices = tf.math.top_k(breed_predictions, k=3)
-                
-                top_breeds = []
-                for i in range(3):
-                    breed_name = class_names[top_k_indices[i]]
-                    confidence = top_k_values[i] * 100
-                    top_breeds.append({"name": breed_name.replace("_", " ").title(), "confidence": f"{confidence:.2f}%"})
-
-                return render_template('index.html',
-                                       is_dog=True,
-                                       predictions=top_breeds,
-                                       uploaded_image=filename)
-            else:
-                # If it's not a dog, return a clear message
-                not_dog_confidence = (1 - gatekeeper_pred) * 100
-                return render_template('index.html',
-                                       is_dog=False,
-                                       prediction_text=f"This doesn't look like a dog.",
-                                       confidence_text=f"({not_dog_confidence:.2f}% sure it's not a dog)",
-                                       uploaded_image=filename)
-
-    return render_template('index.html')
-
-
-@app.route('/uploads/<filename>')
-def uploaded_file(filename):
-    """Serves the uploaded file to be displayed on the webpage."""
-    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
-
-
-if __name__ == '__main__':
-    app.run(debug=True)
+import os
+import numpy as np
+import tensorflow as tf
+from flask import Flask, request, render_template, send_from_directory
+from werkzeug.utils import secure_filename
+from PIL import UnidentifiedImageError # To catch corrupted image errors
+
+# --- 1. CONFIGURATION ---
+app = Flask(__name__)
+UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
+app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+
+# Define paths for the .h5 models
+BREED_MODEL_PATH = 'dog_breed_project_model.h5'      # <-- UPDATED FILENAME
+GATEKEEPER_MODEL_PATH = 'gatekeeper_model.h5'        # <-- UPDATED FILENAME
+CLASS_NAMES_PATH = 'class_names.txt'
+IMG_SIZE = 224
+
+# --- 2. LOAD MODELS AND CLASS NAMES (ONCE AT STARTUP) ---
+print("--- Loading models and class names... ---")
+try:
+    # Load the expert model for breed classification
+    breed_model = tf.keras.models.load_model(BREED_MODEL_PATH)
+    print("✅ Dog Breed (Expert) model loaded.")
+
+    # Load the gatekeeper model for dog vs. not-dog classification
+    gatekeeper_model = tf.keras.models.load_model(GATEKEEPER_MODEL_PATH)
+    print("✅ Gatekeeper (Dog vs. Not-Dog) model loaded.")
+
+    # Load class names
+    with open(CLASS_NAMES_PATH, 'r') as f:
+        class_names = [line.strip() for line in f.readlines()]
+    print(f"✅ Class names loaded. Found {len(class_names)} classes.")
+
+except Exception as e:
+    print(f"❌ Error loading models: {e}")
+    # Define dummy models so the app doesn't crash on startup
+    breed_model = None
+    gatekeeper_model = None
+
+print("--- Setup complete ---")
+
+
+# --- 3. IMAGE PREPROCESSING FUNCTION (CONSISTENT WITH TRAINING) ---
+def preprocess_uploaded_image(filepath, img_size):
+    """
+    Loads, decodes, and preprocesses an uploaded image for both models.
+    This function handles different file types, grayscale images, and aspect ratios.
+    """
+    try:
+        # Read the file and decode it as a 3-channel (RGB) image
+        img = tf.io.read_file(filepath)
+        img = tf.image.decode_image(img, channels=3, expand_animations=False)
+
+        # Pad to a square aspect ratio without distortion
+        img = tf.image.resize_with_pad(img, img_size, img_size)
+        
+        # Expand dimensions to create a batch of 1
+        img_batch = tf.expand_dims(img, 0)
+
+        # Preprocess for each model's specific requirements
+        gatekeeper_input = tf.keras.applications.mobilenet_v2.preprocess_input(tf.identity(img_batch))
+        breed_model_input = tf.keras.applications.resnet_v2.preprocess_input(tf.identity(img_batch))
+        
+        return gatekeeper_input, breed_model_input
+
+    except (UnidentifiedImageError, tf.errors.InvalidArgumentError):
+        # Handle cases where the file is not a valid image
+        return None, None
+    except Exception as e:
+        print(f"An unexpected error occurred during preprocessing: {e}")
+        return None, None
+
+
+# --- 4. FLASK ROUTES ---
+@app.route('/', methods=['GET', 'POST'])
+def index():
+    if request.method == 'POST':
+        if not all([breed_model, gatekeeper_model]):
+             return render_template('index.html', error="Models are not loaded. Please check the server logs.")
+
+        if 'file' not in request.files:
+            return render_template('index.html', error="No file part in the request.")
+        
+        file = request.files['file']
+        if file.filename == '':
+            return render_template('index.html', error="No file selected.")
+
+        if file:
+            filename = secure_filename(file.filename)
+            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+            file.save(filepath)
+
+            # --- PREDICTION PIPELINE ---
+            gatekeeper_img, breed_img = preprocess_uploaded_image(filepath, IMG_SIZE)
+
+            if gatekeeper_img is None:
+                return render_template('index.html', error="Invalid or corrupted image file. Please try another.")
+
+            # Step 1: Use the Gatekeeper to check if the image is a dog
+            gatekeeper_pred = gatekeeper_model.predict(gatekeeper_img)[0][0]
+            
+            if gatekeeper_pred > 0.5:
+                # -----------------------------------------------------------------
+                # START: UPGRADED LOGIC WITH TEST-TIME AUGMENTATION (TTA)
+                # -----------------------------------------------------------------
+                
+                # Create 4 augmented versions of the image (0, 90, 180, 270 degrees)
+                images_to_predict = [
+                    breed_img,
+                    tf.image.rot90(breed_img, k=1), # 90 degrees
+                    tf.image.rot90(breed_img, k=2), # 180 degrees (upside down)
+                    tf.image.rot90(breed_img, k=3)  # 270 degrees
+                ]
+                
+                # Stack the images into a single batch
+                tta_batch = tf.concat(images_to_predict, axis=0)
+                
+                # Get predictions for all 4 images in one go
+                batch_predictions = breed_model.predict(tta_batch)
+                
+                # Average the predictions to get the final, robust result
+                breed_predictions = tf.reduce_mean(batch_predictions, axis=0)
+                
+                # -----------------------------------------------------------------
+                # END: UPGRADED LOGIC WITH TEST-TIME AUGMENTATION (TTA)
+                # -----------------------------------------------------------------
+
+                # Get top 3 predictions from the averaged result
+                top_k_values, top_k_indices = tf.math.top_k(breed_predictions, k=3)
+                
+                top_breeds = []
+                for i in range(3):
+                    breed_name = class_names[top_k_indices[i]]
+                    confidence = top_k_values[i] * 100
+                    top_breeds.append({"name": breed_name.replace("_", " ").title(), "confidence": f"{confidence:.2f}%"})
+
+                return render_template('index.html',
+                                       is_dog=True,
+                                       predictions=top_breeds,
+                                       uploaded_image=filename)
+            else:
+                # If it's not a dog, return a clear message
+                not_dog_confidence = (1 - gatekeeper_pred) * 100
+                return render_template('index.html',
+                                       is_dog=False,
+                                       prediction_text=f"This doesn't look like a dog.",
+                                       confidence_text=f"({not_dog_confidence:.2f}% sure it's not a dog)",
+                                       uploaded_image=filename)
+
+    return render_template('index.html')
+
+
+@app.route('/uploads/<filename>')
+def uploaded_file(filename):
+    """Serves the uploaded file to be displayed on the webpage."""
+    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
+
+
+if __name__ == '__main__':
+    app.run(host="0.0.0.0", port=7860)