Update app.py

app.py

@@ -1,97 +1,71 @@
 import gradio as gr
 import tensorflow as tf
-from PIL import Image
 import numpy as np
-import os
-import logging
-import traceback
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Debug file system
-logger.info("Current directory contents: %s", os.listdir())
-logger.info("Model exists: %s", os.path.exists("model.keras"))
-
-class ImageClassifier:
-    def __init__(self, model_path):
-        try:
-            logger.info("Loading model...")
-            self.model = tf.keras.models.load_model(model_path)
-            logger.info("Model loaded successfully")
-            self.model.summary(print_fn=logger.info)
-        except Exception as e:
-            logger.error("Model loading failed", exc_info=True)
-            raise
-
-        # Update these based on your model
-        self.input_size = (224, 224)
-        self.class_names = ['class1', 'class2', 'class3']
+from PIL import Image
+import json
 
-    def preprocess_image(self, image):
-        try:
-            logger.info("Preprocessing image")
-            image = image.resize(self.input_size)
-            image_array = np.array(image)
-            image_array = image_array / 255.0
-            image_array = np.expand_dims(image_array, axis=0)
-            logger.info("Image preprocessed - shape: %s", image_array.shape)
-            return image_array
-        except Exception as e:
-            logger.error("Preprocessing failed", exc_info=True)
-            raise
+# Load the model (replace 'your_model.keras' with your actual filename)
+model = tf.keras.models.load_model('model.keras')
 
-    def predict(self, image):
-        try:
-            processed_image = self.preprocess_image(image)
-            logger.info("Making prediction...")
-            predictions = self.model.predict(processed_image)
-            logger.info("Raw predictions: %s", predictions)
-            predicted_class = np.argmax(predictions[0])
-            confidence = np.max(predictions[0])
-            return {
-                'class': self.class_names[predicted_class],
-                'confidence': float(confidence),
-                'all_predictions': predictions.tolist()
-            }
-        except Exception as e:
-            logger.error("Prediction failed", exc_info=True)
-            raise
+# Load class labels
+with open("class_labels.json", "r") as f:
+    class_labels = json.load(f)
 
-try:
-    classifier = ImageClassifier("model.keras")
-except Exception as e:
-    logger.error("Failed to initialize classifier", exc_info=True)
-    raise
+def preprocess_image(image):
+    """
+    Preprocess the uploaded image to match model requirements
+    """
+    # Convert PIL image to numpy array
+    img_array = np.array(image)
+    
+    # Resize image to match model input size (256x256 for your model)
+    img_resized = tf.image.resize(img_array, [256, 256])
+    
+    # Normalize pixel values (adjust based on your model's training)
+    img_normalized = tf.cast(img_resized, tf.float32) / 255.0
+    
+    # Add batch dimension
+    img_batch = tf.expand_dims(img_normalized, 0)
+    
+    return img_batch
 
 def classify_image(image):
+    """
+    Classify the uploaded image and return predictions
+    """
     try:
-        logger.info("Received image type: %s", type(image))
-        if isinstance(image, np.ndarray):
-            logger.info("Image shape: %s", image.shape)
-            pil_image = Image.fromarray(image)
-        else:
-            pil_image = image
+        # Preprocess the image
+        processed_image = preprocess_image(image)
+        
+        # Make prediction
+        predictions = model.predict(processed_image)
+        
+        # Get the predicted class probabilities
+        probabilities = tf.nn.softmax(predictions[0]).numpy()
+        
+        # Create a dictionary of class labels and their probabilities
+        results = {}
+        for i, prob in enumerate(probabilities):
+            if i < len(class_labels):
+                results[class_labels[i]] = float(prob)
+        
+        # Sort by probability (highest first)
+        sorted_results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
+        
+        return sorted_results
         
-        logger.info("Converted to PIL format")
-        result = classifier.predict(pil_image)
-        logger.info("Prediction result: %s", result)
-        return result
     except Exception as e:
-        logger.error("Classification error", exc_info=True)
-        return f"Error: {str(e)}"
+        return {"Error": f"Classification failed: {str(e)}"}
 
+# Create Gradio interface
 interface = gr.Interface(
     fn=classify_image,
-    inputs=gr.Image(label="Upload an image"),
-    outputs=[
-        gr.Label(label="Predicted Class"),
-        gr.Number(label="Confidence Score"),
-        gr.JSON(label="All Predictions")
-    ],
-    title="Image Classification",
-    allow_flagging="never"
+    inputs=gr.Image(type="pil", label="Upload an image"),
+    outputs=gr.Label(num_top_classes=5, label="Predictions"),
+    title="Image Classification Model",
+    description="Upload an image to classify it using the trained model",
+    examples=["example1.jpg", "example2.jpg"]  # Add example images if available
 )
 
-interface.launch(debug=True, server_name="0.0.0.0", server_port=7860)
+if __name__ == "__main__":
+    interface.launch()