Add files via upload

hough.py

@@ -1,63 +1,140 @@
-import sys, os.path, cv2, numpy as np
-
-
-def gradient_img(img: np.ndarray) -> np.ndarray:
-    hor_grad = (img[1:, :] - img[:-1, :])[:, :-1]
-    ver_grad = (img[:, 1:] - img[:, :-1])[:-1:, :]
-    magnitude = np.sqrt(hor_grad ** 2 + ver_grad ** 2)
-    return magnitude
-
-
-def hough_transform(
-        img: np.ndarray, theta: float, rho: float
-) -> (np.ndarray, list, list):
-    pass  # insert your code here
-
-
-def get_lines(
-        ht_map: np.ndarray, n_lines: int,
-        thetas: list, rhos: list,
-        min_delta_rho: float, min_delta_theta: float
-) -> list:
-    pass  # insert your code here
-
-
-def main():
-    assert len(sys.argv) == 9
-    src_path, dst_ht_path, dst_lines_path, theta, rho, \
-        n_lines, min_delta_rho, min_delta_theta = sys.argv[1:]
-
-    theta = float(theta)
-    assert theta > 0.0
-
-    rho = float(rho)
-    assert rho > 0.0
-
-    n_lines = int(n_lines)
-    assert n_lines > 0
-
-    min_delta_rho = float(min_delta_rho)
-    assert min_delta_rho > 0.0
-
-    min_delta_theta = float(min_delta_theta)
-    assert min_delta_theta > 0.0
-
-    assert os.path.exists(src_path)
-    img = cv2.imread(src_path, cv2.IMREAD_GRAYSCALE)
-    assert img is not None
-    gradient = gradient_img(img.astype(float))
-
-    ht_map, thetas, rhos = hough_transform(gradient, theta, rho)
-    cv2.imwrite(dst_ht_path, ht_map)
-
-    lines = get_lines(
-        ht_map, n_lines, thetas, rhos, min_delta_rho, min_delta_theta
-    )
-
-    with open(dst_lines_path, 'w') as fout:
-        for line in lines:
-            fout.write(f'{line[0]:.3f}, {line[1]:.3f}\n')
-
-
-if __name__ == '__main__':
-    main()
+import sys, os.path, cv2, math, numpy as np
+
+
+def gradient_img(img: np.ndarray) -> np.ndarray:
+    hor_grad = (img[1:, :] - img[:-1, :])[:, :-1]
+    ver_grad = (img[:, 1:] - img[:, :-1])[:-1:, :]
+    magnitude = np.sqrt(hor_grad ** 2 + ver_grad ** 2)
+    return magnitude
+
+
+def hough_transform(
+        img: np.ndarray, theta: float, rho: float
+) -> (np.ndarray, list, list):
+    hw = img.shape
+
+    # строим необходимые массивы и записываем их размеры
+    rhoSize = int(round(hw[1]/rho) * round(hw[0]/rho))
+    rho2 = round(hw[1]/rho)
+    thetaSize = int(round(180 / theta))
+    rhos = np.zeros(rhoSize)
+    thetas = np.zeros(rhoSize)
+
+
+    ht_map = np.zeros((rhoSize, thetaSize))
+    for i in range(rhoSize):
+        maxi, maxj, max = 0, 0, 0
+        for k in range(int(rho * rho)):
+            #ищем максимальный градиент в нашей окрестности и записываем его данные
+            if(max < img[int(i // rho2 + k // rho)][int(i % rho2 + k % rho)]):
+                maxi = i // rho2 + k // rho
+                maxj = i % rho2 + k % rho
+                max = img[int(i // rho2 + k // rho)][int(i % rho2 + k % rho)]
+        rhos[i] = int(round(math.sqrt(maxi**2 + maxj**2)))
+        thetas[i] = np.arctan(float(maxi)/maxj)
+
+        for j in range(thetaSize):
+            #записываем данные в таблицу среды Хафа
+            A = (math.sin((theta/180) * j * math.pi))
+            B = (math.cos((theta / 180) * j * math.pi))
+            C = (maxj * B - maxi * A) * B
+            ht_map[i][j] = abs(C)/(math.sqrt(A**2 + B**2))
+    return(ht_map, thetas, rhos)
+
+def get_lines(
+        ht_map: np.ndarray, n_lines: int,
+        thetas: list, rhos: list,
+        min_delta_rho: float, min_delta_theta: float,theta
+) -> list:
+
+    out = np.zeros((n_lines, 2))
+    for i in range(ht_map.shape[1]):
+        R = 0
+        KEKlist = np.zeros((ht_map.shape[0], 2))
+        for j in range(ht_map.shape[0]):
+            S = 0
+            for k in range(j):
+                if abs(ht_map[j][i] - KEKlist[k][0]) < min_delta_rho:
+                    KEKlist[k][1] += 1
+                    S += 1
+            if(S == 0):
+                KEKlist[R][1] = 1
+                KEKlist[R][0] = ht_map[j][i]
+                R += 1
+
+    r = n_lines
+    while r > 0:
+        peepo = 0
+        mi = 0
+        mj = 0
+        max = 0
+        maxC = 0
+        for i in range(ht_map.shape[1]):
+            f = 0
+            while f < ht_map.shape[0] and (KEKlist[f][1] != 0 or KEKlist[f][0] != 0):
+                if(KEKlist[f][1] > max):
+                    max = KEKlist[f][1]
+                    maxC = KEKlist[f][0]
+                    mi = i
+                    mj = f
+                f += 1
+        for i in range(ht_map.shape[0]):
+            for j in range(ht_map.shape[1]):
+                if abs(ht_map[i][j] - maxC) < min_delta_rho and max > 0:
+                    A = (math.sin((theta / 180) * mi * math.pi))
+                    B = (math.cos((theta / 180) * mi * math.pi))
+                    C = (rhos[i] * math.cos(thetas[i]) * B - rhos[i] * math.sin(thetas[i]) * A) * B
+                    out[n_lines - r][0] = math.tan(((float(mi) * theta)/180) * math.pi)
+                    out[n_lines - r][1] = C
+                    peepo += 1
+                    break
+            if peepo > 0:
+                break
+        KEKlist[mj][1] = 0
+        r -= 1
+    return out
+
+
+def main():
+    assert len(sys.argv) == 9
+    src_path, dst_ht_path, dst_lines_path, theta, rho, \
+        n_lines, min_delta_rho, min_delta_theta = sys.argv[1:]
+
+    theta = float(theta)
+    assert theta > 0.0
+
+    rho = float(rho)
+    assert rho > 0.0
+
+    n_lines = int(n_lines)
+    assert n_lines > 0
+
+    min_delta_rho = float(min_delta_rho)
+    assert min_delta_rho > 0.0
+
+    min_delta_theta = float(min_delta_theta)
+    assert min_delta_theta > 0.0
+
+    assert os.path.exists(src_path)
+    img = cv2.imread(src_path, cv2.IMREAD_GRAYSCALE)
+
+    assert img is not None
+    gradient = gradient_img(img.astype(float))
+
+
+    ht_map, thetas, rhos = hough_transform(gradient, theta, rho)
+    cv2.imwrite(dst_ht_path, ht_map)
+
+
+
+    lines = get_lines(
+        ht_map, n_lines, thetas, rhos, min_delta_rho, min_delta_theta, theta
+    )
+
+    with open(dst_lines_path, 'w') as fout:
+        for line in lines:
+            fout.write(f'{line[0]:.3f}, {line[1]:.3f}\n')
+
+
+if __name__ == '__main__':
+    main()