src/helpers/template_extraction/script.py

@@ -1,53 +1,14 @@
-""" 
-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
-
-Auteurs: 
-Tom Selier
-Arne van Iterson
-"""
 import cv2
 import numpy as np
 import sys
 import os
+import math as m
 
-## Parameters ##
-
-# [%] Schaal output in percentages 
-SCHAAL = 25         
-
-
-# [mm] 120 op calib a4, 180 op template
-CALIB_AFSTAND = 120
-
-# [mm] design foutje, 10 op a4, 15 op template 
-CORRECTIE = 15
-
-# [mm] mm om af te snijden
-OFFSET = 20         
-
-
-# Toon plaatjes (forceert pauze)
-PLAATJES = False
-
-# Wacht na elke plaatje op toets
-PAUZE = False
-
-# Sla plaatjes op
-SAVE = False        
-
-# extra info
-VERBOSE = True
+# command:
+# python template.py input_folder output_folder
 
+# Returns the correct corner based on IDs
 def getCorner(id):
-    """
-    Returns the correct corner based on IDs
-    """
     if(id == 0):
         return 2
     if(id == 1):
@@ -57,25 +18,8 @@ def getCorner(id):
     if(id == 3):
         return 1
 
-def findMax(array, index):
-    """
-    Vind max voor een rare list van tuples
-    """
-    max_val = int(0)
-    for element in array:
-        if element[index] > max_val:
-            max_val = element[index]
-    return max_val
-
-def findMin(array, index):
-    """
-    Vind min voor een rare list van tuples
-    """
-    min_val = int(10e6)
-    for element in array:
-        if element[index] < min_val:
-            min_val = element[index]
-    return min_val
+if(len(sys.argv) != 3):
+    print("Wrong amount of arguments")
 
 input_folder = sys.argv[1]
 output_folder = sys.argv[2]
@@ -88,42 +32,36 @@ if(not os.path.isdir(output_folder)):
     print("Output folder not found")
     exit()
 
+scale_percentage = 25
 
-### ARUCO PARAMETERS ###
+# Aruco params
 detector_params = cv2.aruco.DetectorParameters()
 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):
-    ## Laad plaatjes in kleur ##
+    # Laad een plaatje
     filename = input_folder + '/' + file
     out_filename = output_folder + '/' + file
-    img = cv2.imread(filename, 1)
+    img = cv2.imread(filename, 0)  
     
-    if VERBOSE:
-        print(file)  
+    # Rescaling 
+    x_scale = int(img.shape[0]/100*scale_percentage)
+    y_scale = int(img.shape[1]/100*scale_percentage)
+    # print("X_scale:", x_scale)
+    # print("Y_scale:", y_scale)
     
-    ## Herschaal de fotos ##
-    (h, w) = img.shape[:2] 
-    x_scale = int(w/100*SCHAAL)
-    y_scale = int(h/100*SCHAAL)
+    # Detecteer markers
     small_img = cv2.resize(img, (x_scale, y_scale))
-    (h_small, w_small) = small_img.shape[:2]
-    if VERBOSE:
-        print("Van %dx%dpx naar %dx%dpx geschaald"%(h,w,h_small,w_small))
-        
-    ## Detecteer aruco markers ##
     (corners, ids, rejected) = detector.detectMarkers(small_img)
     new_image = cv2.aruco.drawDetectedMarkers(
         small_img, 
         corners, 
         ids)
-    assert len(ids) == 4, "Not enough markers found in " + str(file)
-    if VERBOSE:
-        print("%d markers gedetecteerd" %len(ids))
     
-    ## Zoek de 'binnenste' hoekjes ##
+    assert len(ids) == 4, "Not enough markers found"
+
+    # Zoek de "binneste" hoekjes
     i = 0
     x = np.zeros(4, dtype=np.int32)
     y = np.zeros(4, dtype=np.int32)
@@ -133,18 +71,15 @@ for file in os.listdir(input_folder):
         cv2.circle(
             new_image, 
             (x[i], y[i]), 
-            radius=3, 
-            color=(0, 254, 254),
+            radius=5, 
+            color=(255, 100, 255),
             thickness=-1)
+        print("x[%d]: %d"%(i, x[i]))
+        print("y[%d]: %d"%(i, y[i]))
+        print()
         i += 1
-    if VERBOSE:
-        print("Gedecteerde hoekjes (recht): ")
-        print("\t x: ", end='')
-        print(x)
-        print("\t y: ", end='')
-        print(y)
 
-    ## Draai de plaatjes recht ##
+    # Draai plaatje om 0 en 1 gelijk te trekken
     x0 = x[np.where(ids == 0)[0][0]]
     x1 = x[np.where(ids == 1)[0][0]]
     x_dif = x0-x1
@@ -153,69 +88,29 @@ for file in os.listdir(input_folder):
     y1 = y[np.where(ids == 1)[0][0]]
     y_dif = y0 - y1
 
-    angle_rad = np.arctan(y_dif/x_dif)
-    angle_deg = angle_rad * 180.0 / np.pi
-    if(VERBOSE):
-        print("%d graden gedraaid"%angle_deg)
+    angle_error_rad = np.arctan(y_dif/x_dif)
+    angle_error_deg = angle_error_rad * 180.0 / m.pi
+    print("Hoek deg", angle_error_deg)
 
-    (cX, cY) = (w_small // 2, h_small // 2)
-    matrix = cv2.getRotationMatrix2D((cX, cY), angle_deg, 1.0)
-    warp_img = cv2.warpAffine(new_image, matrix, (w_small, h_small))
+    # Calibreer afstand tov afstand tussen 0 en 1
 
-    ## Herdetecteer de binnenste hoekjes op de gedraaide plaatjes ##
-    mask = cv2.inRange(warp_img, (0, 253, 253), (0, 255, 255))
-    warped_corners = []
-    for i in range(4):
-        coordinate = cv2.minMaxLoc(mask)[3]
-        warped_corners.append(coordinate)
-        cv2.circle(mask, warped_corners[i], 10, (0, 0, 0), -1)
 
-    if VERBOSE:
-        print("Gedecteerde hoekjes (gedraaid): ")
-        print("\t",  warped_corners)
+    # Teken rode balken
 
-    ## Teken rode balken voor controle ##
-    # feature voor ooit, nogal lastig
 
-    ## Calibratie van afstand met betrekking tot pixels ##
-    x_min = w_small
-    x_max = 0
-    for corner in warped_corners:
-        # vind x waarden bovenaan
-        (x, y) = corner
-        if y < (h_small // 2):
-            if x < x_min:
-                x_min = x
-            else:
-                x_max = x
-    x_diff = x_max - x_min
-    calib_scale = x_diff / CALIB_AFSTAND
-    if VERBOSE:
-        print("Berekende schaal:  %d pixels/mm"%calib_scale)
+    # Crop
+    offset = 180
+    crop_img = new_image[
+        np.min(x+offset):np.max(x-offset),
+        np.min(y+offset):np.max(y-offset)
+    ]
 
-    ## Crop de plaatjes ##
-    top_left_x = findMin(warped_corners, 0)
-    top_left_y = findMin(warped_corners, 1)
-    bot_right_x = findMax(warped_corners, 0)
-    bot_right_y = findMax(warped_corners, 1)
-
-    offset = int((OFFSET+CORRECTIE)*calib_scale)
-    crop_img = warp_img[
-        top_left_y+offset:bot_right_y-offset, 
-        top_left_x+offset:bot_right_x-offset]
-
-    ## Check voor rode balken ##
-    # zie boven
-
-    ## Output ##
-    if PLAATJES:
-        cv2.imshow('window', warp_img)
+    # Check voor rode balken
+    cv2.imshow('window', new_image)
     cv2.imshow('cropped', crop_img)
-    if SAVE:
-        cv2.imwrite(out_filename, crop_img) 
-    if PAUZE or PLAATJES:
-        cv2.waitKey(0)
-    if VERBOSE:
-        print("=============================================")
 
-cv2.destroyAllWindows()
+    # Keur af of sla op
+    # cv2.imwrite(out_filename, new_image) 
+
+    cv2.waitKey(0)
+cv2.destroyAllWindows