Merge pull request #28 from MarjanAsgari/main

fixing bbox issue. closes #24
NRCan · Oct 24, 2024 · 2bbe63d · 2bbe63d
2 parents ff4fed8 + c365de6
commit 2bbe63d
Show file tree

Hide file tree

Showing 45 changed files with 20 additions and 10 deletions.
diff --git a/.dockerignore b/.dockerignore
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
diff --git a/.gitignore b/.gitignore
diff --git a/Dockerfile b/Dockerfile
diff --git a/LICENSE b/LICENSE
diff --git a/MANIFEST.in b/MANIFEST.in
diff --git a/README.md b/README.md
diff --git a/VERSION b/VERSION
diff --git a/examples/run_geo_inference_with_scripted_model.ipynb b/examples/run_geo_inference_with_scripted_model.ipynb
diff --git a/examples/script_model.py b/examples/script_model.py
diff --git a/examples/test_small.tif b/examples/test_small.tif
diff --git a/geo_inference/__init__.py b/geo_inference/__init__.py
diff --git a/geo_inference/config/__init__.py b/geo_inference/config/__init__.py
diff --git a/geo_inference/config/log_config.yaml b/geo_inference/config/log_config.yaml
diff --git a/geo_inference/config/logging_config.py b/geo_inference/config/logging_config.py
diff --git a/geo_inference/config/sample.yaml b/geo_inference/config/sample.yaml
diff --git a/geo_inference/geo_dask.py b/geo_inference/geo_dask.py
diff --git a/geo_inference/geo_inference.py b/geo_inference/geo_inference.py
@@ -298,11 +298,21 @@ async def async_run_inference(self,
                     top=bbox[3],
                     transform=self.raster_meta["transform"],
                 )
-                col_off, row_off = roi_window.col_off, roi_window.row_off
-                width, height = roi_window.width, roi_window.height
+                self.bbox_transform = from_origin(bbox[0], 
+                    bbox[3],  
+                    self.raster_meta["transform"].a, 
+                    self.raster_meta["transform"].e if self.raster_meta["transform"].e > 0 else -1 * self.raster_meta["transform"].e
+                )
+                col_off, row_off = int(self.roi_window.col_off), int(self.roi_window.row_off)
+                width, height = int(self.roi_window.width), int(self.roi_window.height)
                 aoi_dask_array = aoi_dask_array[
                     :, row_off : row_off + height, col_off : col_off + width
                 ]
+                self.raster_meta.update({
+                    'transform': self.bbox_transform,
+                    'width': aoi_dask_array.shape[2],
+                    'height': aoi_dask_array.shape[1]
+                })
             self.original_shape = aoi_dask_array.shape
             # Pad the array to make dimensions multiples of the patch size
             pad_height = (
@@ -353,7 +363,7 @@ async def async_run_inference(self,
                 pbar.register()
                 import rioxarray
                 logger.info("Inference is running:")
-                aoi_dask_array = xr.DataArray(aoi_dask_array[: self.original_shape[1], : self.original_shape[2]], dims=("y", "x"), attrs= self.json if self.json is not None else xarray_profile_info(self.raster))
+                aoi_dask_array = xr.DataArray(aoi_dask_array[: self.original_shape[1], : self.original_shape[2]], dims=("y", "x"), attrs= self.json if self.json is not None else xarray_profile_info(self.raster_meta))
                 aoi_dask_array.rio.to_raster(mask_path, tiled=True, lock=threading.Lock())
 
             total_time = time.time() - start_time

diff --git a/geo_inference/utils/__init__.py b/geo_inference/utils/__init__.py
diff --git a/geo_inference/utils/geo.py b/geo_inference/utils/geo.py
diff --git a/geo_inference/utils/geo_transforms.py b/geo_inference/utils/geo_transforms.py
diff --git a/geo_inference/utils/helpers.py b/geo_inference/utils/helpers.py
@@ -277,7 +277,7 @@ def select_model_device(gpu_id: int, multi_gpu: bool, device: str="cpu"):
 
 
 def xarray_profile_info(
-    raster,
+    raster_meta
 ):
     """
     Save mask to file.
@@ -286,14 +286,14 @@ def xarray_profile_info(
     Returns:
         None
     """
-    driver = 'GTiff' if raster.driver == 'VRT' else raster.driver
+    raster_meta["driver"] = 'GTiff' if raster_meta["driver"] == 'VRT' else raster_meta["driver"]
     profile_kwargs = {
-        'crs': raster.crs.to_string(),  # Coordinate Reference System, using src.crs.to_string() to get a string representation
-        'transform': raster.transform,  # Affine transformation matrix
+        'crs': raster_meta["crs"],  # Coordinate Reference System, using src.crs.to_string() to get a string representation
+        'transform': raster_meta['transform'],  # Affine transformation matrix
         'count': 1,  # Number of bands
-        'width': raster.width,  # Width of the raster
-        'height': raster.height,  # Height of the raster
-        'driver': driver,  # Raster format driver
+        'width': raster_meta['width'],  # Width of the raster
+        'height': raster_meta['height'],  # Height of the raster
+        'driver': raster_meta["driver"],  # Raster format driver
         'dtype': "uint8",  # Data type (use dtype directly if it's a valid format for xarray)
         'BIGTIFF': 'YES',  # BigTIFF option
         'compress': 'lzw'  # Compression type

diff --git a/geo_inference/utils/polygon.py b/geo_inference/utils/polygon.py
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/requirements.txt b/requirements.txt
diff --git a/tests/__init__.py b/tests/__init__.py
diff --git a/tests/data/0.tif b/tests/data/0.tif
diff --git a/tests/data/0_coco.json b/tests/data/0_coco.json
diff --git a/tests/data/0_corrupt.tif b/tests/data/0_corrupt.tif
diff --git a/tests/data/0_mask.tif b/tests/data/0_mask.tif
diff --git a/tests/data/0_polygons.geojson b/tests/data/0_polygons.geojson
diff --git a/tests/data/0_yolo.csv b/tests/data/0_yolo.csv
diff --git a/tests/data/aff_gdf_result.csv b/tests/data/aff_gdf_result.csv
diff --git a/tests/data/gj_to_px_result.geojson b/tests/data/gj_to_px_result.geojson
diff --git a/tests/data/inference/test_model/test_model.pt b/tests/data/inference/test_model/test_model.pt
diff --git a/tests/data/sample.csv b/tests/data/sample.csv
diff --git a/tests/data/sample.geojson b/tests/data/sample.geojson
diff --git a/tests/data/sample.yaml b/tests/data/sample.yaml
diff --git a/tests/test_geo_dask.py b/tests/test_geo_dask.py
diff --git a/tests/test_geo_inference.py b/tests/test_geo_inference.py
diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
diff --git a/tests/utils/test_geo.py b/tests/utils/test_geo.py
diff --git a/tests/utils/test_geo_transforms.py b/tests/utils/test_geo_transforms.py
diff --git a/tests/utils/test_helpers.py b/tests/utils/test_helpers.py
diff --git a/tests/utils/test_polygon.py b/tests/utils/test_polygon.py