| #!/usr/bin/env python3 | |
| import cv2 | |
| import numpy as np | |
| import os | |
| from joblib import load | |
| class SVMModel: | |
| def __init__(self): | |
| path = os.getenv("SVM_MODEL_PATH", "/home/user/app/model_classification/svm_model.joblib") | |
| self.model = load(path) | |
| def classify_image( | |
| self, | |
| image_bytes: bytes, | |
| image_size=(128, 128) | |
| ) -> int: | |
| img = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR) | |
| if img is None: | |
| # If image fails to load, default to "irrelevant" or handle differently | |
| return 0 | |
| img = cv2.resize(img, image_size) | |
| x = img.flatten().reshape(1, -1) | |
| pred = self.model.predict(x)[0] | |
| return pred | |
| if __name__ == "__main__": | |
| model = load_svm_model("/home/user/app/model_classification/svm_model.joblib") | |
| result = classify_image("test.jpg", model) | |
| print("Classification result:", result) |