This commit is contained in:
nuknal
2024-05-29 10:20:21 +08:00
parent 5eb5ae6a72
commit e15ae9247b
17 changed files with 276 additions and 45 deletions

204
producer/output.go Normal file
View File

@@ -0,0 +1,204 @@
package imageproc
import (
"bufio"
"fmt"
"os"
"path/filepath"
"strings"
"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.Cols()
height := pan.Rows()
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]
band.SetColorInterp(godal.CIGray)
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,
[]godal.ColorInterp{godal.CIBlue, godal.CIGreen, godal.CIRed, godal.CIAlpha},
tiffFile)
}
func (r *Registrator) SavePansharpenedToGDALGTiff(tiffFile string) error {
ihsImage := PansharpenIHS(r.PanImage, r.rgbirImage)
return SaveBGRToGDALGTiff(ihsImage,
3, 1.25,
[]godal.ColorInterp{godal.CIRed, godal.CIGreen, godal.CIBlue},
tiffFile)
}
func SaveBGRToGDALGTiff(bgr gocv.Mat, bands int, resolution float64, colorInterps []godal.ColorInterp, 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()
ds.SetMetadata("NBITS", "16")
setGeoTransform(ds, 0, 0, float64(width), float64(height), resolution)
for b := 0; b < bands; b++ {
band := ds.Bands()[b]
band.SetColorInterp(colorInterps[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 BGR 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
}
func (r *Registrator) SaveRegisteredMssToRaw(raw string) error {
f, err := os.OpenFile(raw, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0777)
if err != nil {
return err
}
var mssData [4][]byte
for i := 0; i < MssBands; i++ {
mssData[i] = r.registeredMssImages[i].ToBytes()
}
width := r.registeredMssImages[0].Cols() * PixelBytes
height := r.registeredMssImages[0].Rows()
log.Println("Writing registered MSS to RAW file...", len(mssData[0])*4)
log.Println("width:", width)
log.Println("height:", height)
w := bufio.NewWriter(f)
for row := 0; row < height; row++ {
w.Write(mssData[0][row*width : (row+1)*width])
w.Write(mssData[1][row*width : (row+1)*width])
w.Write(mssData[2][row*width : (row+1)*width])
w.Write(mssData[3][row*width : (row+1)*width])
}
hdr := EnviHdr{
Samples: 4 * width / PixelBytes,
Lines: height,
Bands: 1,
}
hdrFile := filepath.Join(filepath.Dir(raw),
fmt.Sprintf("%s.HDR", strings.TrimSuffix(filepath.Base(raw), filepath.Ext(raw))),
)
hdr.Write(hdrFile)
return nil
}