import tensorflow as tf
import numpy as np
from PIL import Image
import json
import os

class MobilePhoneUsageDetector:
    def __init__(self, model_path):
        """Initialize the model"""
        self.model = tf.keras.models.load_model(model_path)
        with open('model_config.json', 'r') as f:
            self.config = json.load(f)
        
        self.img_size = tuple(self.config['input_shape'][:2])
        self.class_names = list(self.config['classes'].keys())
    
    def preprocess_image(self, image):
        """Preprocess image for model prediction"""
        if isinstance(image, str):
            image = Image.open(image).convert('RGB')
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # Resize image
        image = image.resize(self.img_size)
        
        # Convert to array and normalize
        image_array = np.array(image) / 255.0
        
        # Add batch dimension
        image_array = np.expand_dims(image_array, axis=0)
        
        return image_array
    
    def predict(self, image):
        """Predict phone usage in image"""
        # Preprocess image
        processed_image = self.preprocess_image(image)
        
        # Make prediction
        prediction = self.model.predict(processed_image)[0][0]
        
        # Get class and confidence
        predicted_class = self.class_names[1] if prediction > 0.5 else self.class_names[0]
        confidence = prediction if prediction > 0.5 else 1 - prediction
        
        return {
            'class': predicted_class,
            'confidence': float(confidence),
            'raw_prediction': float(prediction)
        }

# Example usage
if __name__ == "__main__":
    # Initialize detector
    detector = MobilePhoneUsageDetector('fine_tuned_phone_detection_model.h5')
    
    # Example prediction
    result = detector.predict('example_image.jpg')
    print(f"Prediction: {result}")