Update app/services/predictor.py

app/services/predictor.py

@@ -1,33 +1,44 @@
+# app/services/predictor.py
 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 typing import Tuple, Optional
+
+from app.configs import (
+    get_classification_model,
+    sigmoid_to_class,
+    CLASSIFIER_URI,
+)
 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]:
-    """Predict image using classification model"""
+
+def classify_image(
+    image_bytes: bytes,
+    model_version: Optional[str] = None,
+) -> Tuple[str, float, str]:
+    """
+    Classify image using the classification model.
     
-    # Check if image is valid
-    if not is_valid_image(file_bytes):
-        raise ValueError("Invalid image")
+    Args:
+        image_bytes: Raw image bytes
+        model_version: Optional model version (for A/B testing)
+    
+    Returns:
+        Tuple of (prediction, confidence, actual_version)
+    """
+    # Validate image
+    if not is_valid_image(image_bytes):
+        raise ValueError("Invalid image format")
 
     # Get model
     model = get_classification_model()
-
     if model is None:
         raise ValueError("Classification model not loaded")
 
-    # Save image first
-    import hashlib
-    image_filename = f"{hashlib.md5(filename.encode()).hexdigest()[:12]}.jpg"
-    image_path = save_image(file_bytes, image_filename)
-
     # Preprocess
-    img_array = preprocess_image(file_bytes, target_size=(224, 224))
+    img_array = preprocess_image(image_bytes, target_size=(224, 224))
 
-    # 🔥 Use SavedModel signature
+    # 🔥 Call SavedModel signature
     result = model(tf.constant(img_array))
     
     if isinstance(result, dict):
@@ -35,11 +46,13 @@ def predict_image(file_bytes: bytes, filename: str) -> Tuple[str, float, str]:
     else:
         predictions = result.numpy()
 
-    confidence = float(np.max(predictions[0]))
-    predicted_class = int(np.argmax(predictions[0]))
-
-    prediction = model_classes.get(predicted_class, "unknown")
+    # Sigmoid output -> binary classification
+    confidence = float(predictions[0][0])
+    prediction = sigmoid_to_class(confidence)
+    
+    # Use provided version or default
+    actual_version = model_version or CLASSIFIER_URI
 
-    print(f"🔍 CLASSIFICATION: {prediction} ({confidence:.3f})")
+    print(f"🔍 CLASSIFICATION: {prediction} (confidence={confidence:.3f})")
 
-    return prediction, confidence, image_path
+    return prediction, confidence, actual_version