Spaces:
Sleeping
Sleeping
| """ | |
| AI Model Inference Module | |
| Responsible for loading TensorFlow Keras models and performing AI/real predictions on facial images. | |
| """ | |
| import os | |
| import numpy as np | |
| from PIL import Image | |
| from typing import Tuple, Optional | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class FaceDetectorModel: | |
| """AI face detection model wrapper class""" | |
| def __init__(self, model_path: str = "model/best_mobilenet_finetuned.keras"): | |
| """ | |
| Initialize and load AI detection model | |
| Args: | |
| model_path: Path to model file | |
| Raises: | |
| RuntimeError: If model loading fails | |
| """ | |
| self.model_path = model_path | |
| self.model = None | |
| self.input_size = (224, 224) # Default input size | |
| self.threshold = 0.5 # Classification threshold | |
| try: | |
| self._load_model() | |
| logger.info(f"Model loaded successfully: {model_path}") | |
| except Exception as e: | |
| error_msg = f"Model loading failed: {e}" | |
| logger.error(error_msg) | |
| raise RuntimeError(error_msg) | |
| def _load_model(self): | |
| """Load TensorFlow Keras model""" | |
| try: | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| # Define custom objects to help Keras recognize special layers | |
| custom_objects = { | |
| 'Add': keras.ops.add, | |
| 'Multiply': keras.ops.multiply, | |
| } | |
| # Attempt to load model | |
| try: | |
| self.model = tf.keras.models.load_model( | |
| self.model_path, | |
| custom_objects=custom_objects, | |
| compile=False | |
| ) | |
| except Exception as e: | |
| # If failed, try using safe_mode | |
| logger.warning(f"Standard loading failed, trying safe mode: {e}") | |
| self.model = tf.keras.models.load_model( | |
| self.model_path, | |
| safe_mode=False, # Disable safety checks | |
| compile=False | |
| ) | |
| # Get input shape from model | |
| if hasattr(self.model, 'input_shape'): | |
| input_shape = self.model.input_shape | |
| if len(input_shape) >= 3: | |
| self.input_size = (input_shape[1], input_shape[2]) | |
| except Exception as e: | |
| raise RuntimeError(f"Cannot load model file {self.model_path}: {e}") | |
| def preprocess_image(self, image_path: str) -> Optional[np.ndarray]: | |
| """ | |
| Preprocess image to match model input requirements | |
| Args: | |
| image_path: Path to image file | |
| Returns: | |
| Preprocessed image array, returns None if failed | |
| """ | |
| try: | |
| # Load image | |
| img = Image.open(image_path).convert('RGB') | |
| # Resize to model input size | |
| img = img.resize(self.input_size) | |
| # Convert to numpy array and normalize | |
| img_array = np.array(img, dtype=np.float32) / 255.0 | |
| # Add batch dimension | |
| img_array = np.expand_dims(img_array, axis=0) | |
| return img_array | |
| except Exception as e: | |
| logger.error(f"Image preprocessing failed {image_path}: {e}") | |
| return None | |
| def predict(self, image_path: str) -> Tuple[str, float]: | |
| """ | |
| Predict whether image is AI-generated or real human | |
| Args: | |
| image_path: Path to image file | |
| Returns: | |
| Tuple[str, float]: (prediction label, confidence) | |
| - Prediction label: "AI" or "Human" | |
| - Confidence: Float between 0.0-1.0 | |
| """ | |
| # Preprocess image | |
| img_array = self.preprocess_image(image_path) | |
| if img_array is None: | |
| # Preprocessing failed, return default value | |
| return "AI", 0.5 | |
| try: | |
| # Model prediction | |
| prediction = self.model.predict(img_array, verbose=0)[0][0] | |
| # Convert to label | |
| # Assumption: model output >threshold is AI-generated, <=threshold is real human | |
| label = "AI" if prediction > self.threshold else "Human" | |
| confidence = float(prediction) if prediction > self.threshold else float(1 - prediction) | |
| return label, confidence | |
| except Exception as e: | |
| logger.error(f"Model prediction failed {image_path}: {e}") | |
| return "AI", 0.5 | |
| def warmup(self): | |
| """ | |
| Warmup model using dummy data to reduce first prediction delay | |
| """ | |
| try: | |
| # Create a dummy image | |
| dummy_img = np.random.rand(1, self.input_size[0], self.input_size[1], 3).astype(np.float32) | |
| self.model.predict(dummy_img, verbose=0) | |
| logger.info("Model warmup complete") | |
| except Exception as e: | |
| logger.warning(f"Model warmup failed: {e}") | |