|
|
|
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: |
|
|
|
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) |
