WIP - Model pipeline

src/terrain_diffusion/pipeline.py

@@ -19,3 +19,38 @@
 - It asks Model Inference to run a named model on a patch.
 - It asks Elevation Encoding to turn the models' base and detail outputs into real elevations.
 """
+import numpy as np
+import itertools
+
+from terrain_diffusion.inference import load_model
+from terrain_diffusion import encoding
+
+class ModelPipeline:
+
+  def generate(self, patch: np.ndarray) -> np.ndarray:
+    return self.blur_patch(patch)
+
+  def blur_patch(self, grid: np.ndarry) -> np.ndarray:
+    ret_grid = np.zeros(grid.shape)
+    for (i,j) in np.ndindex(grid.shape):
+
+      top = max(0, i-1)
+      bottom = min(grid.shape[0] - 1, i+1)
+      left = max(0, j-1)
+      right = min(grid.shape[1] - 1, j+1)
+      all_indices = set(itertools.product(range(top, bottom + 1), range(left, right+1)))
+      valid_indices = [(x,y) for (x,y) in all_indices if (x, y) != (i, j)]
+
+      ret_grid[i][j] = sum(grid[x][y] for (x,y) in valid_indices) / len(valid_indices)
+    return ret_grid
+
+
+  def clean_patch(patch: np.ndarray) -> np.ndarray:
+    core, decoder = load_model("core"), load_model("decoder")
+
+    core_output = core.predict(patch)
+    decoder_output = decoder.predict(core_output.latent_map)
+
+    #TODO: update functions and output once elevation encoding is complete
+    some_output = encoding.some_function(core_output.low_res_grid, decoder_output.full_res_grid)
+    return some_output

tests/test_pipeline.py

@@ -0,0 +1,26 @@
+"""
+Tests ModelPipeline generation functions
+"""
+
+import numpy as np
+import pytest
+from terrain_diffusion.pipeline import *
+from terrain_diffusion.inference import PATCH_SIZE
+
+class TestModelPipeline:
+    @pytest.fixture
+    def initial_pipeline(self) -> ModelPipeline:
+        return ModelPipeline()
+
+    @pytest.fixture
+    def initial_patch(self) -> np.ndarray:
+        return np.ones(PATCH_SIZE)
+
+    def test_generate_size(self, initial_pipeline, initial_patch):
+        output = initial_pipeline.generate(initial_patch)
+        assert output.shape == initial_patch.shape, "input and output sizes do not match"
+
+    def test_generate_deterministic(self, initial_pipeline, initial_patch):
+        output_1 = initial_pipeline.generate(initial_patch)
+        output_2 = initial_pipeline.generate(initial_patch)
+        assert output_1 == output_2, "outputs are not the same" #TODO: implement __eq__ for elevation encoding output or smth