added folder parsing

src/experiments/template_extraction/script.py

@@ -3,14 +3,11 @@ Script to extract 10cm x 10cm images from the template
 Accounts for (minor) rotation.  
 
 command:
-python template.py input_folder output_folder 
-
-volgorde maakt uit! heb niet die fancy library gebruikt
+python template.py input_directory
 
 known bugs:
-- plaatjes worden soms de verkeerde kant op gedraaid, ze staan wel recht
-    maar 90 graden gedraaid
 - plaatjes zijn niet persé exact 10x10 cm als perspectief niet precies goed is
+- gekke crash bij herextractie van datapunten, zie todo
 
 Auteurs: 
 Tom Selier
@@ -24,7 +21,7 @@ import os
 ## Parameters ##
 
 # [%] Schaal output in percentages 
-SCHAAL = 25         
+SCHAAL = 100         
 
 
 # [mm] 120 op calib a4, 180 op template
@@ -38,13 +35,13 @@ OFFSET = 20
 
 
 # Toon plaatjes (forceert pauze)
-PLAATJES = True
+PLAATJES = False
 
 # Wacht na elke plaatje op toets
 PAUZE = False
 
 # Sla plaatjes op
-SAVE = False        
+SAVE = True        
 
 # extra info
 VERBOSE = True
@@ -82,30 +79,44 @@ def findMin(array, index):
             min_val = element[index]
     return min_val
 
-input_folder = sys.argv[1]
-output_folder = sys.argv[2]
+input_directory = sys.argv[1]
 
-if(not os.path.isdir(input_folder)):
-    print("Input folder not found")
+if(not os.path.isdir(input_directory)):
+    print(input_directory)
+    print("Input directory not found")
     exit()
 
-if(not os.path.isdir(output_folder)):
-    print("Output folder not found")
-    exit()
-
-
 ### ARUCO PARAMETERS ###
 detector_params = cv2.aruco.DetectorParameters()
+detector_params.adaptiveThreshWinSizeMax = 150 # Takes longer, but better chance of finding markers
 dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_ARUCO_ORIGINAL)  
 detector = cv2.aruco.ArucoDetector(dictionary, detector_params)
 
 ### IMAGE CONVERSIE ###
-for file in os.listdir(input_folder):
+for folder in os.listdir(input_directory):
+    if folder.endswith('.py'):
+        print("Skipping", folder)
+        continue
+    
+    if folder.endswith('.jpg'):
+        continue
+
+    if SAVE:
+        path = os.path.join(input_directory, folder)
+        if (path).endswith("_out"):
+            continue
+        out_folder = path  + "_out"
+        if (os.path.exists(out_folder) == False):
+            if VERBOSE:
+                print("Created folder: " + out_folder)
+            os.mkdir(out_folder)
+
+    for file in os.listdir(os.path.join(input_directory, folder)):
         ## Laad plaatjes in kleur ##
-    filename = input_folder + '/' + file
-    out_filename = output_folder + '/' + file
+        filename = input_directory + '/' + folder + '/' + file
         img = cv2.imread(filename, 1)
 
+        assert img is not None, "Image path is wrong " + filename
         if VERBOSE:
             print(file)  
         
@@ -113,7 +124,10 @@ for file in os.listdir(input_folder):
         (h, w) = img.shape[:2] 
         x_scale = int(w/100*SCHAAL)
         y_scale = int(h/100*SCHAAL)
+        if SCHAAL != 100:
             small_img = cv2.resize(img, (x_scale, y_scale))
+        else:
+            small_img = img
         (h_small, w_small) = small_img.shape[:2]
         if VERBOSE:
             print("Van %dx%dpx naar %dx%dpx geschaald"%(h,w,h_small,w_small))
@@ -124,7 +138,16 @@ for file in os.listdir(input_folder):
             small_img, 
             corners, 
             ids)
-    assert len(ids) == 4, "Not enough markers found in " + str(file)
+        
+        if ids is None:
+            print("Skipping: ", filename)
+            print("=============================================")
+            continue
+        if len(ids) != 4:
+            print("Skipping: ", filename)
+            print("=============================================")
+            continue 
+
         if VERBOSE:
             print("%d markers gedetecteerd" %len(ids))
 
@@ -138,7 +161,7 @@ for file in os.listdir(input_folder):
             cv2.circle(
                 new_image, 
                 (x[i], y[i]), 
-            radius=3, 
+                radius=5, 
                 color=(0, 254, 254),
                 thickness=-1)
             i += 1
@@ -151,7 +174,7 @@ for file in os.listdir(input_folder):
 
         ## Draai de plaatjes recht ##
         x0 = x[np.where(ids == 0)[0][0]]
-    x1 = x[np.where(ids == 1)[0][0]]
+        x1 = x[np.where(ids == 1)[0][0]] # TODO: Fix whatevers happening here
         x_dif = x0-x1
 
         y0 = y[np.where(ids == 0)[0][0]]
@@ -211,7 +234,9 @@ for file in os.listdir(input_folder):
             cv2.imshow('window', warp_img)
             cv2.imshow('cropped', crop_img)
         if SAVE:
-        cv2.imwrite(out_filename, crop_img) 
+            out_path = out_folder + "/" + file
+            print("Saving to:", out_path)
+            cv2.imwrite(out_path, crop_img) 
         if PAUZE or PLAATJES:
             cv2.waitKey(0)
         if VERBOSE: