Add dependencies

app.py

@@ -1,12 +1,8 @@
-import json
 import os
 import gradio as gr
 import plotly.graph_objects as go
 
-import sys
-
-sys.path.append(os.path.join(os.path.dirname(__file__), "..", "SatSeg"))
-from satseg.geo_tools import (
+from geo_tools import (
     shapefile_to_latlong,
     shapefile_to_grid_indices,
     points_to_shapefile,

geo_tools.py

@@ -0,0 +1,105 @@
+import cv2
+import fiona
+from pyproj import Transformer
+
+import numpy as np
+from tqdm import tqdm
+
+from grid_optim import get_optimal_grid
+
+
+def shapefile_to_latlong(file_path: str):
+    c = fiona.open(file_path)
+    contours = []
+
+    transformer = Transformer.from_crs(c.crs, 4326, always_xy=True)
+
+    for poly in c:
+        coords = poly["geometry"].coordinates
+        for coord_set in coords:
+            contours.append(
+                np.array(
+                    list(
+                        transformer.itransform(
+                            coord_set[0] if len(coord_set) == 1 else coord_set
+                        )
+                    )
+                )
+            )
+
+    return contours
+
+
+def shapefile_to_grid_indices(
+    file_path: str, side_len_m: float = 100, meters_per_px: float = 10
+):
+    c = fiona.open(file_path, "r")
+    all_indices = []
+    side_len = side_len_m / meters_per_px
+
+    for poly in c:
+        coords = poly["geometry"].coordinates
+        for coord_set in tqdm(coords):
+            contour = np.array(coord_set[0] if len(coord_set) == 1 else coord_set)
+
+            cmin = contour.min(axis=0)
+            contour -= cmin
+            cmax = int(contour.max() / meters_per_px)
+            contour = contour // meters_per_px
+
+            if cmax < side_len:
+                continue
+
+            mask = np.zeros((cmax, cmax), dtype="uint8")
+            mask = cv2.drawContours(
+                mask,
+                [contour.reshape((-1, 1, 2)).astype(np.int32)],
+                -1,
+                (255),
+                thickness=cv2.FILLED,
+            )
+
+            indices = np.array(get_optimal_grid(mask, side_len=side_len)[0])
+            indices = indices * meters_per_px + cmin
+
+            all_indices += indices.tolist()
+
+    transformer = Transformer.from_crs(c.crs, 4326, always_xy=True)
+    all_indices = list(transformer.itransform(all_indices))
+
+    return np.array(all_indices)
+
+
+def points_to_shapefile(points: np.ndarray, file_path: str):
+    schema = {"geometry": "Point", "properties": [("ID", "int")]}
+
+    # open a fiona object
+    pointShp = fiona.open(
+        file_path,
+        mode="w",
+        driver="ESRI Shapefile",
+        schema=schema,
+        crs="EPSG:4326",
+    )
+    # iterate over each row in the dataframe and save record
+    for i, (x, y) in enumerate(points):
+        rowDict = {
+            "geometry": {"type": "Point", "coordinates": (x, y)},
+            "properties": {"ID": i},
+        }
+        pointShp.write(rowDict)
+    # close fiona object
+    pointShp.close()
+
+
+def get_cached_grid_indices(file_path: str):
+    c = fiona.open(file_path, "r")
+    all_indices = []
+    all_labels = []
+
+    for poly in c:
+        all_indices.append(poly["geometry"].coordinates)
+        if "Name" in poly["properties"]:
+            all_labels.append(poly["properties"]["Name"])
+
+    return np.array(all_indices), all_labels

grid_optim.py

@@ -0,0 +1,113 @@
+import numpy as np
+import math
+from tqdm import tqdm
+
+
+def meters_to_latitude(m_dist: float):
+    return 0.00001 / 1.11 * m_dist
+
+
+def degrees_to_radians(angle_deg: float):
+    return angle_deg / 180 * math.pi
+
+
+def radians_to_degrees(angle_deg: float):
+    return angle_deg / math.pi * 180
+
+
+def plot_line(offset: tuple, start: tuple, indices: list, size: tuple):
+    while (start[0] < 0 or start[1] < 0) and (
+        start[0] < size[0] and start[1] < size[1]
+    ):
+        start = (round(start[0] + offset[0]), round(start[1] + offset[1]))
+    cur_h = start
+
+    while cur_h[0] < size[0] and cur_h[1] < size[1]:
+        indices.append(cur_h)
+        cur_h = (round(cur_h[0] + offset[0]), round(cur_h[1] + offset[1]))
+    return start
+
+
+def get_grid_score(
+    angle: float, side_len: float, seg_result: np.ndarray, threshold: float = 0.5
+):
+    h, w = seg_result.shape
+
+    offset = (int(side_len * math.sin(angle)), int(side_len * math.cos(angle)))
+    s_offset_1 = (
+        int(side_len * math.sin(math.pi / 3 + angle)),
+        int(side_len * math.cos(math.pi / 3 + angle)),
+    )
+    s_offset_2 = (
+        int(side_len * math.sin(math.pi / 3 - angle)),
+        -int(side_len * math.cos(math.pi / 3 - angle)),
+    )
+
+    cur = (0, 0)
+    i = 0
+    indices = []
+
+    # Lower triangle
+    while cur[0] < h and cur[1] < w:
+        cur = plot_line(offset, cur, indices, (h, w))
+        s_offset = s_offset_1 if i % 2 else s_offset_2
+        cur = (round(cur[0] + s_offset[0]), round(cur[1] + s_offset[1]))
+        i += 1
+
+    # Upper triangle
+    cur = (0, 0)
+    for j in range(i + 3):
+        cur = plot_line(offset, cur, indices, (h, w))
+        s_offset = s_offset_1 if j % 2 else s_offset_2
+        cur = (round(cur[0] - s_offset[0]), round(cur[1] - s_offset[1]))
+        j += 1
+
+    indices = np.array(indices)
+    seg_result = (seg_result > threshold) * seg_result
+
+    max_score = 0
+    best_start = tuple()
+    best_indices = None
+    for s_1 in np.arange(0, side_len, side_len // 5):
+        for s_2 in np.arange(0, side_len, side_len // 5):
+            indices_new = indices + np.array([s_1, s_2])
+            indices_new = indices_new[
+                np.logical_and(indices_new[:, 0] < h, indices_new[:, 1] < w)
+            ].astype(int)
+
+            result = np.zeros((h, w))
+            result[indices_new[:, 0], indices_new[:, 1]] = seg_result[
+                indices_new[:, 0], indices_new[:, 1]
+            ]
+
+            score = result.sum()
+            if score > max_score:
+                max_score = score
+                best_start = (s_1, s_2)
+                x, y = np.where(result > threshold)
+                best_indices = np.array([x, y]).T
+
+    # plt.imshow(grid)
+    # plt.show()
+    # plt.imshow(result)
+    # plt.show()
+    return max_score, best_start, np.array([best_indices[:, 1], best_indices[:, 0]]).T
+
+
+def get_optimal_grid(mask: np.ndarray, side_len: float, display_progress: bool = False):
+    angle_range = (0, degrees_to_radians(60))
+    angle_res = degrees_to_radians(10)
+
+    max_score = 0
+    max_config = tuple()
+    best_indices = None
+    for angle in tqdm(
+        np.arange(angle_range[0], angle_range[1], angle_res), disable=True
+    ):
+        score, start, indices = get_grid_score(angle, side_len, mask, 0.8)
+        if score > max_score:
+            max_score = score
+            max_config = (angle, side_len, start)
+            best_indices = indices
+
+    return best_indices, max_config

requirements.txt

@@ -0,0 +1,2 @@
+plotly==5.13.1
+pyproj==3.5.0