update demo script

Author: foamliu

demo.py

@@ -2,9 +2,11 @@
 import cv2 as cv
 import numpy as np
 import scipy.io
-from utils import load_model
+from utils import load_model, draw_str
 import argparse
 import keras.backend as K
+import os
+import random
 
 
 if __name__ == '__main__':
@@ -15,22 +17,27 @@
     class_names = cars_meta['class_names']  # shape=(1, 196)
     class_names = np.transpose(class_names)
 
-    ap = argparse.ArgumentParser()
-    ap.add_argument("-i", "--image", help="path to the image file")
-    args = vars(ap.parse_args())
-
-    filename = args["image"]
-    if filename is None:
-        filename = 'images/samples/07647.jpg'
-
-    bgr_img = cv.imread(filename)
-    rgb_img = cv.cvtColor(bgr_img, cv.COLOR_BGR2RGB)
-    rgb_img = np.expand_dims(rgb_img, 0)
-    preds = model.predict(rgb_img)
-    prob = np.max(preds)
-    class_id = np.argmax(preds)
-    print('class_name: ' + str(class_names[class_id][0][0]))
-    print('prob: ' + str(prob))
+    test_path = 'cars_test/'
+    test_images = [f for f in os.listdir(test_path) if
+                   os.path.isfile(os.path.join(test_path, f)) and f.endswith('.jpg')]
+    samples = random.sample(test_images, 10)
+
+    for i in range(len(samples)):
+        image_name = samples[i]
+        filename = os.path.join('data/test', image_name)
+        print('Start processing image: {}'.format(filename))
+
+        orig_img = cv.imread(os.path.join(test_path, image_name))
+        bgr_img = cv.imread(filename)
+        rgb_img = cv.cvtColor(bgr_img, cv.COLOR_BGR2RGB)
+        rgb_img = np.expand_dims(rgb_img, 0)
+        preds = model.predict(rgb_img)
+        prob = np.max(preds)
+        class_id = np.argmax(preds)
+        text = ('Predict: {}, prob: {}'.format(class_names[class_id][0][0], prob))
+        draw_str(orig_img, (20, 20), text)
+        cv.imwrite('images/{}_out.png'.format(i), orig_img)
+
 
     K.clear_session()
 

images/0_out.png

@@ -1,3 +1,5 @@
+import cv2 as cv
+
 from resnet_152 import resnet152_model
 
 
@@ -8,4 +10,10 @@ def load_model():
     num_classes = 196
     model = resnet152_model(img_height, img_width, num_channels, num_classes)
     model.load_weights(model_weights_path, by_name=True)
-    return model
+    return model
+
+
+def draw_str(dst, target, s):
+    x, y = target
+    cv.putText(dst, s, (x + 1, y + 1), cv.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 0), thickness=2, lineType=cv.LINE_AA)
+    cv.putText(dst, s, (x, y), cv.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255), lineType=cv.LINE_AA)