File size: 930 Bytes
b273357 a6a7c69 0aedcf3 a6a7c69 b273357 a6a7c69 b273357 a6a7c69 b273357 a6a7c69 |
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 |
#!/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_2.joblib")
result = classify_image("test.jpg", model)
print("Classification result:", result) |