File size: 1,767 Bytes
a207b60 fc858b4 a207b60 fc858b4 a207b60 fc858b4 a207b60 f6526ef fc858b4 a207b60 fc858b4 dbe218e a207b60 dbe218e a207b60 b915ef3 | 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 50 51 52 53 54 55 56 57 58 59 60 | import os
import pickle
import cv2
import pandas as pd
import numpy as np
from utils.utils import extract_features_from_image
def run_inference(TEST_IMAGE_PATH, pipeline_model, SUBMISSION_CSV_SAVE_PATH):
test_images = sorted(os.listdir(TEST_IMAGE_PATH))
image_feature_list = []
for test_image in test_images:
path_to_image = os.path.join(TEST_IMAGE_PATH, test_image)
image = cv2.imread(path_to_image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
features = extract_features_from_image(image)
image_feature_list.append(features)
features_multiclass = np.array(image_feature_list)
multiclass_predictions = pipeline_model.predict(features_multiclass)
df_predictions = pd.DataFrame(columns=["file_name", "category_id"])
for i in range(len(test_images)):
file_name = test_images[i]
new_row = pd.DataFrame({
"file_name": file_name,
"category_id": multiclass_predictions[i]
}, index=[0])
df_predictions = pd.concat([df_predictions, new_row], ignore_index=True)
df_predictions.to_csv(SUBMISSION_CSV_SAVE_PATH, index=False)
print(f"Saved predictions to: {SUBMISSION_CSV_SAVE_PATH}")
if __name__ == "__main__":
current_directory = os.path.dirname(os.path.abspath(__file__))
TEST_IMAGE_PATH = "/tmp/data/test_images"
MODEL_NAME = "multiclass_model.pkl"
MODEL_PATH = os.path.join(current_directory, MODEL_NAME)
k = 100
SUBMISSION_CSV_SAVE_PATH = os.path.join(current_directory, "submission.csv")
# load the model
with open(MODEL_PATH, 'rb') as file:
svm_model = pickle.load(file)
run_inference(TEST_IMAGE_PATH, svm_model, SUBMISSION_CSV_SAVE_PATH) |