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

class WasteClassifier:
    def __init__(self, model_path, labels_path):
        # Load the TFLite model
        self.interpreter = tf.lite.Interpreter(model_path=model_path)
        self.interpreter.allocate_tensors()
        
        # Get input and output details
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()
        
        # Load labels
        with open(labels_path, 'r') as f:
            self.labels = [line.strip().split(':')[1] for line in f.readlines()]
    
    def preprocess_image(self, image_path):
        img = Image.open(image_path).convert('RGB')
        img = img.resize((224, 224))
        img_array = np.array(img).astype(np.float32) / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        return img_array
    
    def predict(self, image_path):
        # Preprocess image
        img_array = self.preprocess_image(image_path)
        
        # Set input tensor
        self.interpreter.set_tensor(self.input_details[0]['index'], img_array)
        
        # Run inference
        self.interpreter.invoke()
        
        # Get output tensor
        output_data = self.interpreter.get_tensor(self.output_details[0]['index'])
        predicted_class = np.argmax(output_data[0])
        confidence = float(np.max(output_data[0]))
        
        return self.labels[predicted_class], confidence

# Example usage
if __name__ == "__main__":
    classifier = WasteClassifier('waste_classification.tflite', 'labels.txt')
    prediction, confidence = classifier.predict('sample_image.jpg')
    print(f"Predicted: {prediction} with {confidence:.2%} confidence")