Spaces:
Sleeping
Sleeping
File size: 1,751 Bytes
a4d412e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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")
|