Update app/services/predictor.py

app/services/predictor.py

@@ -1,34 +1,45 @@
 import numpy as np
 from typing import Tuple
+import tensorflow as tf
 from app.configs import get_classification_model, model_classes
 from app.storage import save_image
 from app.core.preprocessing import preprocess_image
 from app.core.validation import is_valid_image
 
 def predict_image(file_bytes: bytes, filename: str) -> Tuple[str, float, str]:
-     # checking if the image is good or no 
+    """Predict image using classification model"""
+    
+    # Check if image is valid
     if not is_valid_image(file_bytes):
         raise ValueError("Invalid image")
 
-    # model loading or ( vonnecting with mlflow )
+    # Get model
     model = get_classification_model()
 
     if model is None:
         raise ValueError("Classification model not loaded")
 
-    image_filename = f"{hash(filename)}.jpg"
+    # Save image first
+    import hashlib
+    image_filename = f"{hashlib.md5(filename.encode()).hexdigest()[:12]}.jpg"
     image_path = save_image(file_bytes, image_filename)
 
-    img_array = preprocess_image(file_bytes)
+    # Preprocess
+    img_array = preprocess_image(file_bytes, target_size=(224, 224))
 
-    # using the model 
-    predictions = model.predict(img_array)
+    # 🔥 Use SavedModel signature
+    result = model(tf.constant(img_array))
+    
+    if isinstance(result, dict):
+        predictions = list(result.values())[0].numpy()
+    else:
+        predictions = result.numpy()
 
-    confidence = float(np.max(predictions))
-    predicted_class = int(np.argmax(predictions))
+    confidence = float(np.max(predictions[0]))
+    predicted_class = int(np.argmax(predictions[0]))
 
     prediction = model_classes.get(predicted_class, "unknown")
 
-    print(f"CLASSIFICATION: {prediction} ({confidence:.3f})")
+    print(f"🔍 CLASSIFICATION: {prediction} ({confidence:.3f})")
 
     return prediction, confidence, image_path