holmeshoo/beans_sorting

1import json
2import numpy as np
3import tensorflow as tf
4from tensorflow import keras
5import cv2
6import sys
7import os
8IMG_HEIGHT = 86
9IMG_WIDTH = 86
10
11
12def predictImg(model, label, img):
13    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
14    img = cv2.resize(img, dsize=(IMG_WIDTH, IMG_HEIGHT))
15    img = img / 255.0
16    img = np.expand_dims(img, axis=0)
17    # モデルによる予測を取得
18    p_result = model.predict(img, verbose=0)
19    print(p_result)
20    result = {}
21    for i in range(len(label)):
22        result[label[i]] = p_result[0][i]
23    return result
24
25
26def predictFile(model, label, img_file_name):
27    # 画像を読み込んで前処理
28    img = cv2.imread(img_file_name)
29    return predictImg(model, label, img)
30
31
32def predictFolder(model, label, folder_path):
33    # フォルダ内のファイル名リストを取得
34    filenames = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]
35    # 各画像に対する予測結果を格納
36    result = []
37    for file in filenames:
38        result.append([file, predictFile(model, label, os.path.join(folder_path, file))])
39    return result
40
41
42def loadModelAndLabel(model_dir, label_label):
43    model = keras.models.load_model(model_dir)
44    with open(label_label, "r") as file:
45        label = json.load(file)
46    return model, label["label"]
47
48
49if __name__ == "__main__":
50    args = sys.argv
51    model_name = args[1]
52    # モデルを読み込み
53    model_f_name = model_name + ".h5"
54    label_f_name = model_name + ".json"
55    model, label = loadModelAndLabel(model_f_name, label_f_name)
56
57    # # フォルダ内の画像に対して予測を行う
58    result = predictFolder(model, label, args[2])
59    print(result)
60    # # 正解クラスが2である画像の数をカウント
61    count = 0
62    for i in result:
63        print("File:", i[0])
64        print("Predictions:", i[1])
65    if 0 == np.argmax(i[1]):
66        count += 1
67    # exapmle console: python3 ./prediction.py "./model/black_beans" "./beens/red/data/ng"
68