Spaces:

yasyn14
/

aurora

Sleeping

aurora / utils /predictor.py

edited docker and main.py

e7b36c2 over 1 year ago

1.75 kB

import numpy as np
from typing import Dict, Any, Union
import tensorflow as tf
import keras
from config import CLASS_NAMES

def predict_skin_condition(img_array: np.ndarray, model: tf.keras.Model) -> Dict[str, Union[str, float]]:
    """
    Predict skin condition from an input image using the provided model.
    
    Args:
        img_array: Input image as numpy array (H, W, C) with RGB channels
        model: Loaded Keras/TensorFlow model for prediction
        
    Returns:
        Dictionary containing predicted condition name and confidence score
    """
    try:
        # Add batch dimension if not present
        if len(img_array.shape) == 3:
            img_array = np.expand_dims(img_array, axis=0)
        
        # Apply EfficientNet preprocessing
        preprocessed_img = keras.applications.efficientnet.preprocess_input(img_array)
        
        # Make prediction (with TF warning suppression)
        with tf.device('/CPU:0'):  # Force CPU prediction for consistent behavior
            pred_probs = model.predict(preprocessed_img, verbose=0)[0]
        
        # Get top prediction
        top_index = np.argmax(pred_probs)
        
        # Ensure index is valid
        if top_index >= len(CLASS_NAMES):
            raise ValueError(f"Predicted index {top_index} exceeds available class names")
        
        return {
            "condition": CLASS_NAMES[top_index],
            "confidence": float(pred_probs[top_index])
        }
    
    except Exception as e:
        # Log error in production systems
        # logger.error(f"Prediction error: {str(e)}")
        
        # Return empty result with error indication
        return {
            "condition": "error",
            "confidence": 0.0
        }

	import numpy as np
	from typing import Dict, Any, Union
	import tensorflow as tf
	import keras
	from config import CLASS_NAMES

	def predict_skin_condition(img_array: np.ndarray, model: tf.keras.Model) -> Dict[str, Union[str, float]]:
	"""
	Predict skin condition from an input image using the provided model.

	Args:
	img_array: Input image as numpy array (H, W, C) with RGB channels
	model: Loaded Keras/TensorFlow model for prediction

	Returns:
	Dictionary containing predicted condition name and confidence score
	"""
	try:
	# Add batch dimension if not present
	if len(img_array.shape) == 3:
	img_array = np.expand_dims(img_array, axis=0)

	# Apply EfficientNet preprocessing
	preprocessed_img = keras.applications.efficientnet.preprocess_input(img_array)

	# Make prediction (with TF warning suppression)
	with tf.device('/CPU:0'): # Force CPU prediction for consistent behavior
	pred_probs = model.predict(preprocessed_img, verbose=0)[0]

	# Get top prediction
	top_index = np.argmax(pred_probs)

	# Ensure index is valid
	if top_index >= len(CLASS_NAMES):
	raise ValueError(f"Predicted index {top_index} exceeds available class names")

	return {
	"condition": CLASS_NAMES[top_index],
	"confidence": float(pred_probs[top_index])
	}

	except Exception as e:
	# Log error in production systems
	# logger.error(f"Prediction error: {str(e)}")

	# Return empty result with error indication
	return {
	"condition": "error",
	"confidence": 0.0
	}