Change label_dir.py to move images to directories named with the label

src/label_image.py

@@ -8,7 +8,7 @@
 image_data = tf.gfile.FastGFile(image_path, 'rb').read()
 
 # Loads label file, strips off carriage return
-label_lines = [line.rstrip() for line 
+label_lines = [line.rstrip() for line
                    in tf.gfile.GFile("/tf_files/retrained_labels.txt")]
 
 # Unpersists graph from file
@@ -20,13 +20,13 @@
 with tf.Session() as sess:
     # Feed the image_data as input to the graph and get first prediction
     softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
-    
+
     predictions = sess.run(softmax_tensor, \
              {'DecodeJpeg/contents:0': image_data})
-    
+
     # Sort to show labels of first prediction in order of confidence
     top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]
-    
+
     for node_id in top_k:
         human_string = label_lines[node_id]
         score = predictions[0][node_id]

src/py/label_dir.py

@@ -1,58 +1,67 @@
+import errno
 import tensorflow as tf
-import sys
+from shutil import move
+from os import listdir, makedirs
+from os.path import isfile, join, isdir
 
-# change this as you see fit
-#image_path = sys.argv[1]
+def mkdir_p(path):
+    try:
+        makedirs(path)
+    except OSError as exc:  # Python >2.5
+        if exc.errno == errno.EEXIST and isdir(path):
+            pass
+        else:
+            raise
 
-# Read in the image_data
-#image_data = tf.gfile.FastGFile(image_path, 'rb').read()
-import os
-import shutil
-from os import listdir
-from os import mkdir
-from shutil import copyfile
-from os.path import isfile, join
-varPath = '/toScan'
-destDir = "/scanned"
-imgFiles = [f for f in listdir(varPath) if isfile(join(varPath, f))]
 
+src_dir = '/toScan'
+dest_dir = '/scanned'
+img_files = [f for f in listdir(src_dir) if isfile(join(src_dir, f))]
+issue_dir = join(dest_dir, '_issues')
+mkdir_p(issue_dir)
 
 # Loads label file, strips off carriage return
-label_lines = [line.rstrip() for line 
-                   in tf.gfile.GFile("/tf_files/retrained_labels.txt")]
+label_lines = [line.rstrip() for line
+                   in tf.gfile.GFile('/tf_files/retrained_labels.txt')]
 
 # Unpersists graph from file
-with tf.gfile.FastGFile("/tf_files/retrained_graph.pb", 'rb') as f:
+with tf.gfile.FastGFile('/tf_files/retrained_graph.pb', 'rb') as f:
     graph_def = tf.GraphDef()
     graph_def.ParseFromString(f.read())
     _ = tf.import_graph_def(graph_def, name='')
 
 with tf.Session() as sess:
     # Feed the image_data as input to the graph and get first prediction
     softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
-    #try:
-    #    shutil.rmtree(destDir)
-    #except:
-    #    None
-    #mkdir ('scanned')
-
-    for imageFile in imgFiles:
-        image_data =  tf.gfile.FastGFile(varPath+"/"+imageFile, 'rb').read()       
-
-        print (varPath+"/"+imageFile)
-        predictions = sess.run(softmax_tensor, \
-                 {'DecodeJpeg/contents:0': image_data})
-
+
+    for image_file in img_files:
+        src_image_path = join(src_dir, image_file)
+        image_data =  tf.gfile.FastGFile(src_image_path, 'rb').read()
+
+        print(src_image_path)
+
+        try:
+
+            predictions = sess.run(softmax_tensor, \
+                     {'DecodeJpeg/contents:0': image_data})
+        except Exception as e:
+            move(srcFilePath, join(issue_dir, imageFile))
+            continue
+
         # Sort to show labels of first prediction in order of confidence
         top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]
         firstElt = top_k[0];
 
-        newFileName = label_lines[firstElt] +"--"+ str(predictions[0][firstElt])[2:7]+".jpg"
-        print(newFileName)
-        copyfile(varPath+"/"+imageFile, destDir+"/"+newFileName)
+        label_dest_folder = join(dest_dir, label_lines[firstElt])
+        mkdir_p(label_dest_folder)
+
+        final_file_path = join(label_dest_folder, image_file)
+
+        print(final_file_path)
+
+        move(src_image_path, final_file_path)
 
         for node_id in top_k:
             human_string = label_lines[node_id]
             score = predictions[0][node_id]
-            #print (node_id)
             print('%s (score = %.5f)' % (human_string, score))