labels = ['can', 'paper_cup', 'paper_box', 'paper_milkbox', 'plastic'] print("[*] Importing packages...") import tensorflow as tf from tensorflow import keras import os import cv2 import pandas as pd import numpy as np print("[*] Loading model...") os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' tf.get_logger().setLevel('ERROR') model = keras.models.load_model('model_without_preprocess_finetuned.h5') def predict(img, weight): prob = model.predict([np.expand_dims(img, 0), np.expand_dims(weight, 0)]) print(prob) result = prob.argmax(-1)[0] return labels[result] print("[*] Warming up model...") img = cv2.imread(os.path.join('test_data', 'can_dew_0_preprocessed.jpg')) print(predict(img, 52.69)) print("[*] Done!") if __name__ == '__main__': def pred(f, dirpath): img = cv2.imread(os.path.join(dirpath, f)) weight = df.loc[f]['weight'] prob = model.predict([np.expand_dims(img, 0), np.expand_dims(weight, 0)]) result = prob.argmax(-1)[0] return labels[result] df = pd.read_csv('test_data/weights_test.csv') df = df.set_index('name') for dirpath, dirnames, filenames in os.walk('test_data'): for f in filenames: if f.endswith('.jpg'): print(f'{f}: {pred(f, dirpath)}')