Merge branch 'main' of https://arnweb.nl/gitea/arne/EV5_Beeldherk_Bomen

src/helpers/template_extraction/script.py

@@ -1,14 +1,53 @@
+""" 
+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
 
-# command:
-# python template.py input_folder output_folder
+## 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
 
-# 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):
@@ -18,8 +57,25 @@ def getCorner(id):
     if(id == 3):
         return 1
 
-if(len(sys.argv) != 3):
-    print("Wrong amount of arguments")
+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
 
 input_folder = sys.argv[1]
 output_folder = sys.argv[2]
@@ -32,36 +88,42 @@ if(not os.path.isdir(output_folder)):
     print("Output folder not found")
     exit()
 
-scale_percentage = 25
 
-# Aruco params
+### ARUCO PARAMETERS ###
 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 een plaatje
+    ## Laad plaatjes in kleur ##
     filename = input_folder + '/' + file
     out_filename = output_folder + '/' + file
-    img = cv2.imread(filename, 0)  
+    img = cv2.imread(filename, 1)
 
-    # 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)
+    if VERBOSE:
+        print(file)  
     
-    # Detecteer markers
+    ## Herschaal de fotos ##
+    (h, w) = img.shape[:2] 
+    x_scale = int(w/100*SCHAAL)
+    y_scale = int(h/100*SCHAAL)
     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))
 
-    assert len(ids) == 4, "Not enough markers found"
-
-    # Zoek de "binneste" hoekjes
+    ## Zoek de 'binnenste' hoekjes ##
     i = 0
     x = np.zeros(4, dtype=np.int32)
     y = np.zeros(4, dtype=np.int32)
@@ -71,15 +133,18 @@ for file in os.listdir(input_folder):
         cv2.circle(
             new_image, 
             (x[i], y[i]), 
-            radius=5, 
-            color=(255, 100, 255),
+            radius=3, 
+            color=(0, 254, 254),
             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 plaatje om 0 en 1 gelijk te trekken
+    ## Draai de plaatjes recht ##
     x0 = x[np.where(ids == 0)[0][0]]
     x1 = x[np.where(ids == 1)[0][0]]
     x_dif = x0-x1
@@ -88,29 +153,69 @@ for file in os.listdir(input_folder):
     y1 = y[np.where(ids == 1)[0][0]]
     y_dif = y0 - y1
 
-    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)
+    angle_rad = np.arctan(y_dif/x_dif)
+    angle_deg = angle_rad * 180.0 / np.pi
+    if(VERBOSE):
+        print("%d graden gedraaid"%angle_deg)
 
-    # Calibreer afstand tov afstand tussen 0 en 1
+    (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))
 
+    ## 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)
 
-    # Teken rode balken
+    if VERBOSE:
+        print("Gedecteerde hoekjes (gedraaid): ")
+        print("\t",  warped_corners)
 
+    ## Teken rode balken voor controle ##
+    # feature voor ooit, nogal lastig
     
-    # Crop
-    offset = 180
-    crop_img = new_image[
-        np.min(x+offset):np.max(x-offset),
-        np.min(y+offset):np.max(y-offset)
-    ]
+    ## 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)
 
-    # Check voor rode balken
-    cv2.imshow('window', new_image)
-    cv2.imshow('cropped', crop_img)
+    ## 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)
 
-    # Keur af of sla op
-    # cv2.imwrite(out_filename, new_image) 
+    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]
 
-    cv2.waitKey(0)
-cv2.destroyAllWindows
+    ## Check voor rode balken ##
+    # zie boven
+
+    ## Output ##
+    if PLAATJES:
+        cv2.imshow('window', warp_img)
+        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()