Update app.py

app.py

@@ -5,23 +5,39 @@ import json
 from tensorflow.keras.preprocessing.image import img_to_array
 from tensorflow.keras.applications.vgg16 import preprocess_input
 from PIL import Image
-import os
+import re
+
+# ---------- Helper Function: Clean Class Names ----------
+def clean_class_name(raw_name):
+    """
+    Converts raw class names like 'apple_red_1' to 'Apple Red'
+    and 'pear_1' to 'Pear'.
+    """
+    # Remove the trailing underscore and number (e.g., '_1', '_2')
+    # Using regex to remove '_' followed by digits at the end
+    cleaned = re.sub(r'_\d+$', '', raw_name)
+    # Replace underscores with spaces
+    cleaned = cleaned.replace('_', ' ')
+    # Capitalize each word
+    cleaned = cleaned.title()
+    return cleaned
 
 # ---------- Load the Model (.keras format) ----------
-# The model file 'fruits_classifier.keras' should be in the same directory
 model_path = 'fruits_classifier.keras'
 model = tf.keras.models.load_model(model_path)
 print("Model loaded successfully!")
 
 # ---------- Load Class Names ----------
-# Load the class index mapping saved from the training generator
 with open('class_indices.json', 'r') as f:
     class_names_dict = json.load(f)
 
-# Convert dictionary to list for easy index-based lookup
-# Example: {"0": "apple", "1": "banana"} -> ["apple", "banana"]
+# Convert dictionary to a list for easy index-based access
 class_names_list = [class_names_dict[str(i)] for i in range(len(class_names_dict))]
-print(f"Total classes loaded: {len(class_names_list)}")
+
+# ---------- Create a list of cleaned class names for display ----------
+# Format: ["Apple Red", "Apple Braeburn", "Cucumber", ...]
+cleaned_class_names = [clean_class_name(name) for name in class_names_list]
+print(f"Total classes loaded: {len(cleaned_class_names)}")
 
 # ---------- Prediction Function ----------
 def predict_image(image):
@@ -29,7 +45,7 @@ def predict_image(image):
     Takes an input image, preprocesses it for VGG16, runs inference,
     and returns the predicted class name and confidence.
     """
-    # Resize image to match model's expected input shape (64x64)
+    # Resize image to match the model's input shape (64x64)
     img = image.resize((64, 64))
     # Convert PIL image to numpy array
     img_array = img_to_array(img)
@@ -38,19 +54,38 @@ def predict_image(image):
     # Apply VGG16-specific preprocessing (scaling and mean subtraction)
     img_array = preprocess_input(img_array)
     
-    # Get model predictions
+    # Run inference
     predictions = model.predict(img_array)
     predicted_index = np.argmax(predictions, axis=-1)[0]
     confidence = np.max(predictions, axis=-1)[0]
     
-    # Get the corresponding class name
-    predicted_class = class_names_list[predicted_index]
+    # Get the raw class name
+    raw_class_name = class_names_list[predicted_index]
+    # Clean the class name for display
+    predicted_class = clean_class_name(raw_class_name)
     confidence_percentage = float(confidence) * 100
     
     return predicted_class, f"{confidence_percentage:.2f}%"
 
+# ---------- Create a formatted list of categories for display ----------
+# Format: "Apple, Apple Braeburn, Apple Crimson Snow, ..."
+categories_list = sorted(cleaned_class_names)  # Sort alphabetically
+categories_text = ", ".join(categories_list)   # Join with commas
+
+# Create the description with categories
+description_text = f"""
+### 🍎 Upload an image of a fruit.
+
+**The model can predict the following {len(categories_list)} categories:**
+
+{', '.join(categories_list)}
+
+---
+
+*Model: VGG16-based Transfer Learning trained on Fruits-360 dataset.*
+"""
+
 # ---------- Gradio Interface Setup ----------
-# Define the Gradio UI
 interface = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil", label="Upload Fruit Image"),
@@ -59,7 +94,7 @@ interface = gr.Interface(
         gr.Textbox(label="📊 Confidence")
     ],
     title="🍎 Fruit Classification Using Transfer Learning",
-    description="Upload an image of a fruit. The model (VGG16-based) will classify it."
+    description=description_text,
 )
 
 # ---------- Launch the App ----------