Add mock rail detection algorithm

mocka/octiva.py

@@ -0,0 +1,72 @@
+import math
+from pathlib import Path
+import cv2 as cv
+import numpy as np
+
+
+def find_parallel(lines) -> []:
+    """Find parallel lines and return their index."""
+    parallel_lines = []
+    # TODO Can probably be done using cross-product/zip
+    for i in range(len(lines)):
+        for j in range(len(lines)):
+            if i == j:
+                # Every line is parallel to itself. We don't want all lines...
+                continue
+            if abs(lines[i][0][1] - lines[j][0][1]) == 0:
+                parallel_lines.append(i)
+    return parallel_lines
+
+
+def rail_finder_algo_one(filename: str) -> None:
+    """Create image with the lines laid over."""
+    # TODO Don't load image twice
+    source_color_image = cv.imread(cv.samples.findFile(filename), cv.IMREAD_COLOR)
+    source_greyscale_image = cv.imread(cv.samples.findFile(filename), cv.IMREAD_GRAYSCALE)
+    edge_detected_image = cv.Canny(source_greyscale_image, 50, 200, None, 3)
+    # edge_detected_image_color_format = cv.cvtColor(edge_detected_image, cv.COLOR_GRAY2BGR)
+    detected_lines = cv.HoughLinesP(edge_detected_image, 1, np.pi / 180, 150, None, 250, 0)
+    # if detected_lines is not None:
+    #     for i in range(0, len(detected_lines)):
+    #         l = detected_lines[i][0]
+    #         cv.line(edge_detected_image_color_format, (l[0], l[1]), (l[2], l[3]), (0, 0, 255), 6, cv.LINE_AA)
+    parallel_line_indices = find_parallel(detected_lines)
+    # Color found lines red
+    for detected_line in detected_lines:
+        x1, y1, x2, y2 = detected_line[0]
+        cv.line(source_color_image, (x1, y1), (x2, y2), (0, 0, 255), 6, cv.LINE_AA)
+    # Color parallel lines green
+    for parallel_line_index in parallel_line_indices:
+        x1, y1, x2, y2 = detected_lines[parallel_line_index][0]
+        cv.line(source_color_image, (x1, y1), (x2, y2), (0, 255, 0), 6, cv.LINE_AA)
+    cv.imshow("Probabilistic Line Transform", source_color_image)
+
+
+def rail_finder_algo_two(filename):
+    source_greyscale_image = cv.imread(cv.samples.findFile(filename), cv.IMREAD_GRAYSCALE)
+    dst = cv.Canny(source_greyscale_image, 50, 200, None, 3)
+    cdst = cv.cvtColor(dst, cv.COLOR_GRAY2BGR)
+    lines = cv.HoughLines(dst, 1, np.pi / 180, 450, None, 0, 0)
+    if lines is not None:
+        for i in range(0, len(lines)):
+            rho = lines[i][0][0]
+            theta = lines[i][0][1]
+            a = math.cos(theta)
+            b = math.sin(theta)
+            x0 = a * rho
+            y0 = b * rho
+            pt1 = (int(x0 + 1000 * (-b)), int(y0 + 1000 * (a)))
+            pt2 = (int(x0 - 1000 * (-b)), int(y0 - 1000 * (a)))
+            cv.line(cdst, pt1, pt2, (0, 0, 255), 3, cv.LINE_AA)
+    cv.imshow("Standard Hough Line", cdst)
+
+
+def main():
+    filename = str(Path(__file__).parent.parent / Path("tests/octiva") / Path("rails_0.jpg"))
+    rail_finder_algo_one(filename)
+    # rail_finder_algo_two(filename)
+    cv.waitKey()
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -7,6 +7,7 @@ pyarrow           ~= 13.0.0
 requests          ~= 2.31.0
 starlette         ~= 0.27.0
 python-magic      ~= 0.4.27
+opencv-python     ~= 4.8.0.76
 uvicorn[standard] ~= 0.23.2
 # Test
 httpx  ~= 0.25.0