Skip to content

Commit

Permalink
Update add block save
Browse files Browse the repository at this point in the history
  • Loading branch information
geoyee committed Dec 23, 2021
1 parent 6cda4f1 commit 3a5692b
Show file tree
Hide file tree
Showing 9 changed files with 150 additions and 250 deletions.
12 changes: 6 additions & 6 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -82,15 +82,15 @@ This work is in progress, at present, the relevant documents are as follows :
- [x] Add note:
- [x] About how to training your data in AI Studio / Local.
- [x] About different model (paper's link).
- [ ] Accelerate:
- [x] Accelerate and reduce memory:
- [x] PaddlePaddle setting.
- [x] Add maximum pixelsize to calculate / using GDAL `translat / warp` to make raster smaller.
- [ ] Block stacking and saving.
- [ ] Add online map tiles support:
- [ ] Extract building on raster in memory.
- [ ] Add vector range selection.

- [x] Block stacking and saving.
- [ ] Test:
- [x] On Windows 10/11.
- [ ] On Linux.
- [ ] On mac OS Big Sur+.

- [ ] Add online map tiles support:
- [ ] Extract building on raster in memory.
- [ ] Add vector range selection.
84 changes: 38 additions & 46 deletions buildseg/buildSeg.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,7 @@ def unload(self):
action)
self.iface.removeToolBarIcon(action)


# Load parameters
def select_params_file(self):
self.param_file = self.dlg.mQfwParams.filePath()
Expand All @@ -216,9 +217,9 @@ def select_params_file(self):
"use_bf16": self.dlg.ccbBF16.isChecked()
}
self.infer_worker.load_model(self.model_file, self.param_file, use_setting)
print("Parameters loaded successfully")
print("Parameters loaded successfully!")
else:
print(f"Parameters loaded unsuccessfully, not find {self.model_file}")
print(f"Parameters loaded unsuccessfully, not find {self.model_file}.")


# Select shapefile save path
Expand All @@ -231,6 +232,7 @@ def simp_state_change(self, state):
self.dlg.lblThreshold.setEnabled(bool(state // 2))
self.dlg.mQgsDoubleSpinBox.setEnabled(bool(state // 2))


# chackbox state
def gpu_state_change(self, state):
if self.dlg.ccbGPU.isChecked():
Expand Down Expand Up @@ -265,13 +267,14 @@ def __display_error(info_txt):
import numpy
import paddle
except ImportError:
__display_error("Please check if `numpy / opencv-python / paddlepaddle` exists in your environment!")
__display_error("Please check if `numpy / opencv-python / paddlepaddle` " + \
"exists in your environment!")
self.first_start = True
return False
# check paddlepaddle's version
vers = paddle.__version__.split(".")
if int(vers[0]) < 2 or int(vers[1]) < 2:
__display_error("Please make sure your paddlepaddle's version is greater than 2.2.0")
__display_error("Please make sure your paddlepaddle's version is greater than 2.2.0.")
self.first_start = True
return False
# global import utils
Expand Down Expand Up @@ -301,12 +304,12 @@ def init_setting(self):
# # quick test in my computer
# self.dlg.cbxScale.setCurrentIndex(5)
# self.dlg.mQfwShape.setFilePath(r"C:\Users\Geoyee\Desktop\dd\test.shp")
# self.dlg.mQfwParams.setFilePath(r"E:\dataFiles\github\buildseg\static_weight\bisenet_v2_512x512\model.pdiparams")
# self.dlg.mQfwParams.setFilePath(
# r"E:\dataFiles\github\buildseg\static_weight\bisenet_v2_512x512\model.pdiparams")


def run(self):
"""Run method that performs all the real work"""

# Create the dialog with elements (after translation) and keep reference
# Only create GUI ONCE in callback, so that it will only load when the plugin is started
if self.first_start == True:
Expand All @@ -315,77 +318,66 @@ def run(self):
self.init_setting() # init all of widget's settings
# check env
check_pass = self.check_python_pip_env()
if check_pass is True:
if check_pass is True: # env ok
self.infer_worker = utils.InferWorker(self.model_file, self.param_file)
# Run the dialog event loop
result = self.dlg.exec_()
# See if OK was pressed
if result:
# Start timing
time_start = time.time()
# Do something useful here - delete the line containing pass and
# substitute with your code.
# Get parameters
grid_size = [int(self.dlg.cbxBlock.currentText())] * 2
overlap = [int(self.dlg.cbxOverlap.currentText())] * 2
scale_rate = float(self.dlg.cbxScale.currentText())
print(f"grid_size is {grid_size}, overlap is {overlap}, scale_rate is {scale_rate}")
print(f"grid_size is {grid_size}, overlap is {overlap}, scale_rate is {scale_rate}.")
# layers = iface.activeLayer() # Get the currently active layer
current_raster_layer = self.dlg.mMapLayerComboBoxR.currentLayer() # Get the selected raster layer
band_list = current_raster_layer.renderer().usesBands() # Band used by the current renderer
current_raster_layer_name = current_raster_layer.source() # Get the raster layer path
# Get the selected raster layer
current_raster_layer = self.dlg.mMapLayerComboBoxR.currentLayer()
# Band used by the current renderer
band_list = current_raster_layer.renderer().usesBands()
# Get the raster layer path
current_raster_layer_name = current_raster_layer.source()
# Add downsample
layer_path = utils.dowm_sample(current_raster_layer_name, scale_rate)
# open raster to get tf and proj
# fn_ras = QgsProject.instance().mapLayersByName(current_raster_layer_name)[0]
# ras_path = str(fn_ras.dataProvider().dataSourceUri())
down_save_path = self.save_shp_path.replace(".shp", "_dowm.tif")
layer_path = utils.dowm_sample(current_raster_layer_name, down_save_path, scale_rate)
print(f"layer_path: {layer_path}")
ras_ds = gdal.Open(layer_path)
geot = ras_ds.GetGeoTransform()
proj = ras_ds.GetProjection()
# proj = layers.crs()
# If this layer is a raster layer
xsize = ras_ds.RasterXSize
ysize = ras_ds.RasterYSize
grid_count, mask_grids = utils.create_grids(ysize, xsize, grid_size, overlap)
ras_ds = None
grid_count = utils.create_grids(ysize, xsize, grid_size, overlap)
number = grid_count[0] * grid_count[1]
# print(f"xsize is {xsize}, ysize is {ysize}, grid_count is {grid_count}") # test
print("Start block processing")
print("Start block processing.")
geoinfo = {"row": ysize, "col": xsize, "geot": geot, "proj": proj}
mask_save_path = self.save_shp_path.replace(".shp", "_mask.tif")
mask = utils.Mask(mask_save_path, geoinfo, grid_size, overlap)
for i in range(grid_count[0]):
for j in range(grid_count[1]):
img = utils.layer2array(layer_path, band_list, i, j, grid_size, overlap)
# cv2.imwrite("C:/Users/Geoyee/Desktop/dd/" + str(i) + "-" + str(j) + ".jpg", img) # test
mask_grids[i][j] = self.infer_worker.infer(img, True)
# cv2.imwrite("C:/Users/Geoyee/Desktop/dd/" + str(i) + "-" + str(j) + ".png", mask_grids[i][j]) # test
print(f"-- {i * grid_count[1] + j + 1}/{number} --")
print("Start Spliting")
mask = utils.splicing_grids(mask_grids, ysize, xsize, grid_size, overlap)
# cv2.imwrite("C:/Users/Geoyee/Desktop/test.png", mask) # test
print("Start to extract the boundary")
# # raster to shapefile used OpenCV
# build_bound = bound2shp(get_polygon(mask),
# get_transform(layers))
# vl = showgeoms([build_bound], "Building boundary", proj=proj)
# if self.save_shp_path is not None:
# QgsVectorFileWriter.writeAsVectorFormat(
# vl, self.save_shp_path, "utf-8",
# driverName="ESRI Shapefile")
# print(f"Save the Shapefile in {self.save_shp_path}")
# # raster to shapefile used GDAL
# mask_grids[i][j] = self.infer_worker.infer(img, True)
mask.write_grid(self.infer_worker.infer(img, True), i, j)
print(f"-- {i * grid_count[1] + j + 1}/{number} --.")
print("Start Spliting.")
# mask = utils.splicing_grids(mask_grids, ysize, xsize, grid_size, overlap)
print("Start to extract the boundary.")
# raster to shapefile used GDAL
is_simp = self.dlg.ccbSimplify.isChecked()
utils.polygonize_raster(mask, self.save_shp_path, proj, geot, display=(not is_simp))
utils.polygonize_raster(mask, self.save_shp_path, proj, geot, \
display=(not is_simp))
if is_simp is True:
simp_save_path = osp.join(osp.dirname(self.save_shp_path), \
osp.basename(self.save_shp_path).replace(".shp", "_simp.shp"))
utils.simplify_polygon(self.save_shp_path, \
simp_save_path, \
simp_save_path = self.save_shp_path.replace(".shp", "_simp.shp")
utils.simplify_polygon(self.save_shp_path, simp_save_path, \
self.dlg.mQgsDoubleSpinBox.value())
iface.addVectorLayer(simp_save_path, "deepbands-simplified", "ogr")
else :
print ('No')
# # Reset model params
# self.infer_worker.reset_model()
time_end = time.time()
iface.messageBar().pushMessage(
f"The whole operation is performed in less than {str(time_end - time_start)} seconds",
f"The whole operation is performed in {str(time_end - time_start)} seconds.",
level=Qgis.Info,
duration=30)
3 changes: 1 addition & 2 deletions buildseg/utils/__init__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
from .infer import InferWorker
from .qgser import showgeoms, get_transform, bound2shp
from .convert import layer2array
from .splicing import create_grids, splicing_grids
from .boundary import get_polygon
from .splicing import create_grids, Mask # splicing_grids
from .shape import polygonize_raster
from .simplify import simplify_polygon, dowm_sample
147 changes: 0 additions & 147 deletions buildseg/utils/boundary.py

This file was deleted.

2 changes: 1 addition & 1 deletion buildseg/utils/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def layer2array(sample_path, band_list, row=None, col=None, grid_size=[512, 512]
array = np.stack(array_list, axis=2)
else:
array = raster_to_uint8(__get_grid(gd, row, col, \
width, height, grid_size, overlap))
width, height, grid_size, overlap))
del gd
return array

Expand Down
13 changes: 9 additions & 4 deletions buildseg/utils/shape.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import os
import os.path as osp
import numpy as np
from qgis.utils import iface

try:
Expand All @@ -12,9 +13,8 @@

def __mask2tif(mask, tmp_path, proj, geot):
row, columns = mask.shape[:2]
dim = 1
driver = gdal.GetDriverByName("GTiff")
dst_ds = driver.Create(tmp_path, columns, row, dim, gdal.GDT_UInt16)
dst_ds = driver.Create(tmp_path, columns, row, 1, gdal.GDT_UInt16)
dst_ds.SetGeoTransform(geot)
dst_ds.SetProjection(proj)
dst_ds.GetRasterBand(1).WriteArray(mask)
Expand All @@ -23,8 +23,12 @@ def __mask2tif(mask, tmp_path, proj, geot):


def polygonize_raster(mask, shp_save_path, proj, geot, rm_tmp=True, display=True):
tmp_path = shp_save_path.replace(".shp", ".tif")
ds = __mask2tif(mask, tmp_path, proj, geot)
if type(mask) is np.ndarray:
tmp_path = shp_save_path.replace(".shp", ".tif")
ds = __mask2tif(mask, tmp_path, proj, geot)
else:
tmp_path = mask.file_name
ds = mask.gdal_data
srcband = ds.GetRasterBand(1)
maskband = srcband.GetMaskBand()
gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
Expand Down Expand Up @@ -52,6 +56,7 @@ def polygonize_raster(mask, shp_save_path, proj, geot, rm_tmp=True, display=True
dst_ds.Destroy()
ds = None
if rm_tmp:
mask.close()
os.remove(tmp_path)
if display:
iface.addVectorLayer(shp_save_path, "deepbands", "ogr")
Expand Down
Loading

0 comments on commit 3a5692b

Please sign in to comment.