initial.py

import os
import onnx

from skimage.io import imread
from skimage.transform import resize

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, confusion_matrix

from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# prepare data
input_dir = './bug-bites'
categories = ['flea','mosquito','tick']

data = []
labels = []

for category_idx, category in enumerate(categories):
    for file in os.listdir(os.path.join(input_dir, category)):
        img_path = os.path.join(input_dir, category, file)
        img = imread(img_path)
        img = resize(img, (30, 30, 3), anti_aliasing=True)
        data.append(img.flatten())
        labels.append(category_idx)

data = np.asarray(data)
labels = np.asarray(labels)

# train / test split
x_train, x_test, y_train, y_test = train_test_split(data, labels, test_size=0.2, shuffle=True, stratify=labels)

# train classifier
classifier = SVC()

parameters = [{'gamma': [0.01, 0.001, 0.0001], 'C': [1, 10, 100, 1000]}]

grid_search = GridSearchCV(classifier, parameters)

grid_search.fit(x_train, y_train)

# test performance
best_estimator = grid_search.best_estimator_

y_prediction = best_estimator.predict(x_test)

score = accuracy_score(y_prediction, y_test)

print(confusion_matrix(y_prediction,y_test))
print(f"{score*100}% correctly id")

#convert the model to ONNX
intial_type = [('float_input', FloatTensorType ([None, x_train.shape[1]]))]
onnx_model=convert_sklearn(best_estimator, initial_types=intial_type, name='SVMClassifier')

#save the model
with open('./model.onnx', 'wb') as f:
    f.write(onnx_model.SerializeToString())