154 lines
3.6 KiB
Go
154 lines
3.6 KiB
Go
package imageproc
|
||
|
||
import (
|
||
"bufio"
|
||
"os"
|
||
|
||
"github.com/airbusgeo/godal"
|
||
log "github.com/sirupsen/logrus"
|
||
"gocv.io/x/gocv"
|
||
)
|
||
|
||
func (r *Registrator) SaveOriginalPanToGDALGTiff(tiffFile string) error {
|
||
return savePanToGDALGTiff(r.PanImage, tiffFile)
|
||
}
|
||
|
||
func (r *Registrator) SaveFilteredPanToGDALGTiff(tiffFile string) error {
|
||
img := PANFilter(r.PanImage)
|
||
img.ConvertTo(&img, gocv.MatTypeCV16U)
|
||
return savePanToGDALGTiff(img, tiffFile)
|
||
}
|
||
|
||
func savePanToGDALGTiff(pan gocv.Mat, tiffFile string) error {
|
||
log.Println("Saving PAN image to TIFF file:", tiffFile)
|
||
|
||
width := pan.Rows()
|
||
height := pan.Cols()
|
||
|
||
ds, err := godal.Create(godal.GTiff, tiffFile, 1, godal.UInt16, width, height)
|
||
if err != nil {
|
||
log.Error("Error creating TIFF file: ", err)
|
||
return err
|
||
}
|
||
defer ds.Close()
|
||
|
||
setGeoTransform(ds, 0, 0, float64(width), float64(height), 1.25)
|
||
ds.SetMetadata("NBITS", "16")
|
||
|
||
// 将通道的数据转换为uint16数组
|
||
data := make([]uint16, width*height)
|
||
for y := 0; y < height; y++ {
|
||
for x := 0; x < width; x++ {
|
||
data[y*width+x] = uint16(pan.GetShortAt(y, x))
|
||
}
|
||
}
|
||
|
||
band := ds.Bands()[0]
|
||
err = band.IO(godal.IOWrite,
|
||
0, 0,
|
||
data,
|
||
width, height,
|
||
godal.PixelSpacing(2),
|
||
godal.LineSpacing(width*2))
|
||
if err != nil {
|
||
log.Error("Failed to write data to band:", err)
|
||
return err
|
||
}
|
||
|
||
log.Info("Saved pan image to ", tiffFile)
|
||
|
||
return nil
|
||
}
|
||
|
||
func (r *Registrator) SaveRegisteredMssToGDALGTiff(tiffFile string) error {
|
||
width := r.MssWidth
|
||
height := r.MssHeight
|
||
|
||
// 创建合并后的图像（RGBIR）
|
||
r.rgbirImage = gocv.NewMatWithSize(height, width, gocv.MatTypeCV16UC4) // 4通道，16位
|
||
|
||
for y := 0; y < height; y++ {
|
||
for x := 0; x < width; x++ {
|
||
blue := r.registeredMssImages[0].GetShortAt(y, x)
|
||
green := r.registeredMssImages[1].GetShortAt(y, x)
|
||
red := r.registeredMssImages[2].GetShortAt(y, x)
|
||
ir := r.registeredMssImages[3].GetShortAt(y, x)
|
||
r.rgbirImage.SetShortAt(y, x*4+0, blue)
|
||
r.rgbirImage.SetShortAt(y, x*4+1, green)
|
||
r.rgbirImage.SetShortAt(y, x*4+2, red)
|
||
r.rgbirImage.SetShortAt(y, x*4+3, ir)
|
||
}
|
||
}
|
||
|
||
return SaveBGRToGDALGTiff(r.rgbirImage, 4, 5, tiffFile)
|
||
}
|
||
|
||
func (r *Registrator) SavePansharpenedToGDALGTiff(tiffFile string) error {
|
||
ihsImage := PansharpenIHS(r.PanImage, r.rgbirImage)
|
||
return SaveBGRToGDALGTiff(ihsImage, 3, 1.25, tiffFile)
|
||
}
|
||
|
||
func SaveBGRToGDALGTiff(bgr gocv.Mat, bands int, resolution float64, tiffFile string) error {
|
||
log.Println("Saving BGR to TIFF file:", tiffFile)
|
||
|
||
width := bgr.Cols()
|
||
height := bgr.Rows()
|
||
|
||
// 创建一个二维切片来存储图像数据
|
||
data := make([][]uint16, bands)
|
||
for i := range data {
|
||
data[i] = make([]uint16, width*height)
|
||
}
|
||
|
||
// 从gocv.Mat中提取数据
|
||
for y := 0; y < height; y++ {
|
||
for x := 0; x < width; x++ {
|
||
for b := 0; b < bands; b++ {
|
||
data[b][y*width+x] = uint16(bgr.GetShortAt(y, x*bands+b))
|
||
}
|
||
}
|
||
}
|
||
|
||
ds, err := godal.Create(godal.GTiff,
|
||
tiffFile,
|
||
bands,
|
||
godal.UInt16,
|
||
width, height)
|
||
if err != nil {
|
||
log.Error("Error creating TIFF file: ", err)
|
||
return err
|
||
}
|
||
defer ds.Close()
|
||
|
||
setGeoTransform(ds, 0, 0, float64(width), float64(height), resolution)
|
||
|
||
for b := 0; b < bands; b++ {
|
||
band := ds.Bands()[b]
|
||
err := band.IO(godal.IOWrite,
|
||
0, 0,
|
||
data[b],
|
||
width, height,
|
||
godal.PixelSpacing(2),
|
||
godal.LineSpacing(width*2))
|
||
if err != nil {
|
||
log.Error("Failed to write data to band:", err)
|
||
return err
|
||
}
|
||
}
|
||
|
||
log.Info("Saved registered mss to ", tiffFile)
|
||
|
||
return nil
|
||
}
|
||
|
||
func (r *Registrator) BytesToRaw(mssData []byte, filePath string) error {
|
||
f, err := os.OpenFile(filePath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0777)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
w := bufio.NewWriter(f)
|
||
w.Write(mssData)
|
||
|
||
return nil
|
||
}
|