Skip to content

Commit

Permalink
gdal fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
@rafaqz
rafaqz committed Oct 23, 2023
1 parent 3183553 commit 0289a05
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 139 deletions.
188 changes: 71 additions & 117 deletions ext/RastersArchGDALExt/gdal_source.jl
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,5 @@
const AG = ArchGDAL

const GDAL_X_ORDER = ForwardOrdered()
const GDAL_Y_ORDER = ReverseOrdered()

const GDAL_LOCUS = Start()

# drivers supporting the gdal Create() method to directly write to disk
Expand Down Expand Up @@ -61,61 +58,29 @@ Write a `Raster` to file using GDAL.
Returns `filename`.
"""
function Base.write(
filename::AbstractString, ::Type{GDALsource}, A::AbstractRaster{T,2};
force=false, verbose=true, kw...
) where T
RA.check_can_write(filename, force)
all(hasdim(A, (X, Y))) || error("Array must have Y and X dims")

correctedA = _maybe_correct_to_write(A)
nbands = 1
_write(filename, correctedA, nbands; kw...)
end
function Base.write(
filename::AbstractString, ::Type{GDALsource}, A::AbstractRaster{T,3};
filename::AbstractString, ::Type{GDALsource}, A::AbstractRaster{T};
force=false, verbose=true, kw...
) where T
RA.check_can_write(filename, force)
all(hasdim(A, (X, Y))) || error("Array must have Y and X dims")
hasdim(A, Band()) || error("Must have a `Band` dimension to write a 3-dimensional array")

correctedA = _maybe_correct_to_write(A)
nbands = size(correctedA, Band())
_write(filename, correctedA, nbands; kw...)
A1 = _maybe_correct_to_write(A)
_create_with_driver(filename, dims(A1), eltype(A1), missingval(A1); _block_template=A1, kw...) do dataset
open(A1; write=true) do O
AG.RasterDataset(dataset) .= parent(O)
end
end
return filename
end

function RA.create(filename, ::Type{GDALsource}, T::Type, dims::DD.DimTuple;
missingval=nothing, metadata=nothing, name=nothing,
driver="", lazy=true,
options=Dict{String,String}(), _block_template=nothing
missingval=nothing, metadata=nothing, name=nothing, lazy=true, kw...
)
x, y = nolookup_to_sampled(DD.dims(dims, (XDim, YDim)))
T = Missings.nonmissingtype(T)

if ismissing(missingval)
missingval = RA._writeable_missing(T)
end

if hasdim(dims, Band)
b = DD.dims(dims, Band)
nbands = length(b)
newdims = (x, y, b)
else
nbands = 1
newdims = (x, y)
end

create_kw = (width=length(x), height=length(y), nbands=nbands, dtype=T)

_create_with_driver(filename, driver, create_kw; _block_template, options) do dataset
_set_dataset_properties!(dataset, A)
missingval = ismissing(missingval) ? RA._writeable_missing(T) : missingval
_create_with_driver(filename, dims, T, missingval; kw...) do _
nothing
end

if hasdim(dims, Band)
return Raster(filename; source=GDALsource, name, lazy)
else
return view(Raster(filename; source=GDALsource, name, lazy), Band(1))
end
return Raster(filename; source=GDALsource, name, lazy, dropband=!hasdim(dims, Band))
end

function RA._open(f, ::Type{GDALsource}, filename::AbstractString; write=false, kw...)
Expand Down Expand Up @@ -156,7 +121,7 @@ function DD.dims(raster::AG.RasterDataset, crs=nothing, mappedcrs=nothing)
catch
GDAL_EMPTY_TRANSFORM
end
@show gt_dims
# @show gt_dims
gt = gt_dims
xsize, ysize = size(raster)
nbands = AG.nraster(raster)
Expand Down Expand Up @@ -186,7 +151,7 @@ function DD.dims(raster::AG.RasterDataset, crs=nothing, mappedcrs=nothing)
xindex = LinRange(xmin, xmax, xsize)
xorder = xstep > 0 ? ForwardOrdered() : ReverseOrdered()

ystep = gt[GDAL_NS_RES] # A negative number
ystep = gt[GDAL_NS_RES] # Usually a negative number
if ystep > 0
ymax = gt[GDAL_TOPLEFT_Y]
ymin = gt[GDAL_TOPLEFT_Y] + ystep * (ysize - 1)
Expand Down Expand Up @@ -269,7 +234,6 @@ function RA.Raster(ds::AG.RasterDataset;
lazy=false,
dropband=false
)
@show dims
args = dims, refdims, name, metadata, missingval
filelist = AG.filelist(ds)
raster = if lazy && length(filelist) > 0
Expand Down Expand Up @@ -316,34 +280,11 @@ function AG.Dataset(f::Function, A::AbstractRaster; kw...)
f(rds.ds)
end
end
function AG.RasterDataset(f::Function, A::AbstractRaster;
filename=nothing, driver="", options, _block_template=A
)
all(hasdim(A, (X, Y))) || throw(ArgumentError("`AbstractRaster` must have both an `X` and `Y` to be converted to an ArchGDAL `Dataset`"))
if ndims(A) === 3
thirddim = otherdims(A, (X, Y))[1]
thirddim isa Band || throw(ArgumentError("ArchGDAL can't handle $(basetypeof(thirddim)) dims - only X, Y, and Band"))
elseif ndims(A) > 3
throw(ArgumentError("ArchGDAL can only accept 2 or 3 dimensional arrays"))
end

A_m = RA._maybe_use_type_missingval(A, GDALsource)
A_p = _maybe_permute_to_gdal(A_m)

kw = (;
width=length(DD.dims(A_p, X)),
height=length(DD.dims(A_p, Y)),
nbands=hasdim(A_p, Band) ? size(A_p, Band()) : 1,
dtype=eltype(A_p),
)
if filename == nothing
filename = ""
end

return _create_with_driver(filename, driver, kw; _block_template) do dataset
_set_dataset_properties!(dataset, dims(A_p), missingval(A_p))
function AG.RasterDataset(f::Function, A::AbstractRaster; filename="", kw...)
A1 = _maybe_correct_to_write(A)
return _create_with_driver(filename, dims(A1), eltype(A1), missingval(A1); _block_template=A1, kw...) do dataset
rds = AG.RasterDataset(dataset)
open(A_p) do a
open(A1) do a
rds .= parent(a)
end
f(rds)
Expand Down Expand Up @@ -376,64 +317,76 @@ _missingval_from_gdal(T, x) = x
_maybe_correct_to_write(A) = _maybe_correct_to_write(lookup(A, X()), A)
_maybe_correct_to_write(lookup, A) = A
function _maybe_correct_to_write(lookup::Union{AbstractSampled,NoLookup}, A)
_maybe_permute_to_gdal(A) |>
a -> RA.nolookup_to_sampled(a)
end

# Write a Raster to disk using GDAL
function _write(filename, A::AbstractRaster, nbands; driver="", kw...)
A = RA._maybe_use_type_missingval(A, GDALsource)
create_kw = (width=size(A, X()), height=size(A, Y()), nbands=nbands, dtype=eltype(A))
_create_with_driver(filename, driver, create_kw; _block_template=A, kw...) do dataset
_set_dataset_properties!(dataset, A)
rds = AG.RasterDataset(dataset)
open(A; write=true) do O
rds .= parent(O)
end
end

return filename
RA._maybe_use_type_missingval(A, GDALsource) |> _maybe_permute_to_gdal
end

_check_driver(filename::Nothing, driver) = "MEM"
function _check_driver(filename::AbstractString, driver)
if isempty(driver)
driver = AG.extensiondriver(filename)
if driver == "COG"
driver = "GTiff"
if isempty(driver)
if isempty(filename)
driver = "MEM"
else
driver = AG.extensiondriver(filename)
if driver == "COG"
driver = "GTiff"
end
end
end
return driver
end

# Handle creating a dataset with any driver,
# applying the function `f` to the created dataset
function _create_with_driver(f, filename, driver, create_kw;
options=Dict{String,String}(), _block_template=nothing
function _create_with_driver(f, filename, dims, T, missingval;
options=Dict{String,String}(), driver="", _block_template=nothing, kw...
)
_gdal_validate(dims)

x, y = map(DD.dims(dims, (XDim, YDim))) do d
maybeshiftlocus(Start(), RA.nolookup_to_sampled(d))
end
newdims = hasdim(dims, Band()) ? (x, y, DD.dims(dims, Band)) : (x, y)
nbands = hasdim(dims, Band) ? length(DD.dims(dims, Band())) : 1

driver = _check_driver(filename, driver)
options_vec = _process_options(driver, options; _block_template)
gdaldriver = driver isa String ? AG.getdriver(driver) : driver

create_kw = (; width=length(x), height=length(y), nbands, dtype=T,)
filename = isnothing(filename) ? "" : filename

if AG.shortname(gdaldriver) in GDAL_DRIVERS_SUPPORTING_CREATE
AG.create(filename; driver=gdaldriver, create_kw..., options=options_vec) do dataset
AG.create(filename; driver=gdaldriver, options=options_vec, create_kw...) do dataset
_set_dataset_properties!(dataset, newdims, missingval)
f(dataset)
end
else
# Create a memory object and copy it to disk, as ArchGDAL.create
# Create a tif and copy it to `filename`, as ArchGDAL.create
# does not support direct creation of ASCII etc. rasters
ArchGDAL.create(""; driver=AG.getdriver("MEM"), create_kw...) do dataset
result = f(dataset)
# This `copy` copies _to disk_
AG.copy(dataset; filename=filename, driver=gdaldriver, options=options_vec) |> AG.destroy
result
tif_options_vec = _process_options("GTiff", Dict{String,String}(); _block_template)
tif_driver = AG.getdriver("GTiff")
tif_name = tempname() * ".tif"
AG.create(tif_name; driver=tif_driver, options=tif_options_vec, create_kw...) do dataset
_set_dataset_properties!(dataset, newdims, missingval)
f(dataset)
target_ds = AG.copy(dataset; filename=filename, driver=gdaldriver, options=options_vec)
AG.destroy(target_ds)
end
end
end

@noinline function _gdal_validate(dims)
all(hasdim(dims, (XDim, YDim))) || throw(ArgumentError("`Raster` must have both an `X` and `Y` to be converted to an ArchGDAL `Dataset`"))
if length(dims) === 3
otherdim = otherdims(dims, (XDim, YDim))[1]
otherdim isa Band || throw(ArgumentError("ArchGDAL can't handle $(basetypeof(thirddim)) dims - only X, Y, and Band"))
elseif !(length(dims) in (2, 3))
throw(ArgumentError("ArchGDAL can only accept 2 or 3 dimensional arrays"))
end
end

# Convert a Dict of options to a Vector{String} for GDAL
function _process_options(driver::String, options::Dict;
_block_template=nothing
)
function _process_options(driver::String, options::Dict; _block_template=nothing)
# Get the GDAL driver object
gdaldriver = AG.getdriver(driver)

Expand Down Expand Up @@ -509,7 +462,8 @@ _set_dataset_properties!(ds::AG.Dataset, A) =
_set_dataset_properties!(ds, dims(A), missingval(A))
function _set_dataset_properties!(dataset::AG.Dataset, dims::Tuple, missingval)
# We cant write mixed Points/Intervals, so default to Intervals if mixed
if any(x -> x isa Intervals, sampling.(dims)) && any(x -> x isa Points, sampling.(dims))
xy = DD.dims(dims, (X, Y))
if any(x -> x isa Intervals, sampling.(xy)) && any(x -> x isa Points, sampling.(xy))
dims = set(dims, X => Intervals, Y => Intervals)
end
# Convert the dimensions to `Projected` if they are `Converted`
Expand All @@ -533,7 +487,7 @@ function _set_dataset_properties!(dataset::AG.Dataset, dims::Tuple, missingval)

# Set the geotransform from the updated lookups
gt = RA.dims2geotransform(x, y)
@show gt
# @show gt
AG.setgeotransform!(dataset, gt)

# Set the missing value/nodataval. This is a little complicated
Expand All @@ -554,7 +508,7 @@ function _set_dataset_properties!(dataset::AG.Dataset, dims::Tuple, missingval)

# Write band labels if they are not Integers.
if hasdim(dims, Band)
bandlookup = DD.lookup(dims, Band)
bandlookup = DD.lookup(DD.dims(dims, Band))
if !(eltype(bandlookup) <: Integer)
for i in eachindex(bandlookup)
AG.getband(dataset, i) do band
Expand Down Expand Up @@ -587,9 +541,9 @@ _maybe_permute_to_gdal(A, dims::Tuple) = A
_maybe_permute_to_gdal(A, dims::Tuple{<:XDim,<:YDim,<:Band}) = permutedims(A, dims)
_maybe_permute_to_gdal(A, dims::Tuple{<:XDim,<:YDim}) = permutedims(A, dims)

_maybe_permute_from_gdal(A, dims::Tuple) = permutedims(A, dims)
_maybe_permute_from_gdal(A, dims::Tuple{<:XDim,<:YDim,<:Band}) = A
_maybe_permute_from_gdal(A, dims::Tuple{<:XDim,<:YDim}) = A
_maybe_restore_from_gdal(A, dims::Tuple) = reorder(permutedims(A, dims), dims)
_maybe_restore_from_gdal(A, dims::Tuple{<:XDim,<:YDim,<:Band}) = reorder(A, dims)
_maybe_restore_from_gdal(A, dims::Tuple{<:XDim,<:YDim}) = reorder(A, dims)

#= Geotranforms ########################################################################
Expand Down
2 changes: 1 addition & 1 deletion ext/RastersArchGDALExt/resample.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ function resample(A::RasterStackOrArray;
elseif size isa Tuple{<:Dimension{Int},<:Dimension{Int}}
map(val, dims(size, (YDim, XDim)))
elseif size isa Tuple{Int,Int}
reverse(size)
dimnum(A, XDim) > dimnum(A, YDim) ? size : reverse(size)
else
throw(ArgumentError("`size` must be a `Int`, or a 2 `Tuple` of `Int` or `Dimension`s wrapping `Int`. Got $size"))
end
Expand Down
27 changes: 6 additions & 21 deletions ext/RastersArchGDALExt/warp.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,41 +15,26 @@ function warp(st::AbstractRasterStack, flags::Dict; filename=nothing, suffix=key
end

function _warp(A::AbstractRaster, flags::Dict; filename=nothing, suffix="", kw...)

@show "start"
A1 = _set_gdalwarp_sampling(A)
# @show sampling.(lookup(A1, (X, Y))) bounds(A1, (X, Y))
filename = RA._maybe_add_suffix(filename, suffix)
flagvect = reduce([flags...]; init=[]) do acc, (key, val)
append!(acc, String[_asflag(key), _stringvect(val)...])
end
tempfile = isnothing(filename) ? nothing : tempname() * ".tif"
warp_kw = isnothing(filename) || filename == "/vsimem/tmp" ? () : (; dest=filename)
@show "here"
out = AG.Dataset(A1; filename=tempfile, kw...) do dataset
A2 = Raster(dataset)
# @show sampling.(lookup(A2, (X, Y))) bounds(A2, (X, Y))
rds = Raster(dataset)
AG.gdalwarp([dataset], flagvect; warp_kw...) do warped
# @show dims(AG.RasterDataset(warped))
@show "before"
# Read the raster lazily, dropping Band if there is none in `A`
raster = Raster(warped; lazy=true, dropband=!hasdim(A, Band()))
@show "after"
# @show sampling.(lookup(raster, (X, Y))) bounds(raster, (X, Y)) bounds(raster, Y)[2]

# Either read the MEM dataset, or get the filename as a FileArray
# And permute the dimensions back to what they were in A
p_raster = _maybe_permute_from_gdal(raster, dims(A))

# Either read the MEM dataset to an Array, or keep a filename base raster lazy
return isnothing(filename) ? read(p_raster) : p_raster
return isnothing(filename) ? read(raster) : raster
end
end
# @show sampling.(lookup(out, (X, Y))) bounds(out, (X, Y))
x = _reset_gdalwarp_sampling(out, A)
# @show sampling.(lookup(x, (X, Y))) bounds(x, (X, Y))
println()
return x
# And permute the dimensions back to what they were in A
out1 = _maybe_restore_from_gdal(out, dims(A))
out2 = _reset_gdalwarp_sampling(out1, A)
return out2
end

_asflag(x) = string(x)[1] == '-' ? x : string("-", x)
Expand Down

0 comments on commit 0289a05

Please sign in to comment.