270 lines
7.2 KiB
Go
270 lines
7.2 KiB
Go
package main
|
||
|
||
import (
|
||
"bufio"
|
||
"fmt"
|
||
"image"
|
||
"os"
|
||
|
||
"github.com/chai2010/tiff"
|
||
log "github.com/sirupsen/logrus"
|
||
"gocv.io/x/gocv"
|
||
)
|
||
|
||
type Registrate interface{}
|
||
|
||
const (
|
||
MssBands = 4
|
||
PixelBytes = 2
|
||
PanWidth = 9344 // 像素宽度
|
||
MssWidth = 2336
|
||
BlockNum = 5
|
||
)
|
||
|
||
type Registrator struct {
|
||
PanImage gocv.Mat
|
||
PanHeight int
|
||
PanWidth int
|
||
|
||
MssImages [4]gocv.Mat
|
||
MssHeight int
|
||
MssWidth int
|
||
registeredImages [4]gocv.Mat
|
||
rgbirImage gocv.Mat
|
||
}
|
||
|
||
func NewRegistrator() *Registrator {
|
||
var r Registrator
|
||
|
||
return &r
|
||
}
|
||
|
||
func (r *Registrator) LoadPanRaw(raw string) error {
|
||
data, err := os.ReadFile(raw)
|
||
if err != nil {
|
||
log.Error("Error reading raw file: ", err)
|
||
return err
|
||
}
|
||
|
||
height := len(data) / (PanWidth * PixelBytes)
|
||
r.PanImage, err = gocv.NewMatFromBytes(height, PanWidth, gocv.MatTypeCV16U, data)
|
||
if err != nil {
|
||
log.Error("Error creating Mat from bytes: ", err)
|
||
return err
|
||
}
|
||
|
||
r.PanHeight = height
|
||
r.PanWidth = PanWidth
|
||
|
||
return nil
|
||
}
|
||
|
||
func (r *Registrator) LoadMssRaw(raw string) error {
|
||
data, err := os.ReadFile(raw)
|
||
if err != nil {
|
||
log.Error("Error reading raw file: ", err)
|
||
return err
|
||
}
|
||
|
||
height := len(data) / (MssWidth * PixelBytes * MssBands)
|
||
mssData := make([][]byte, MssBands)
|
||
for h := 0; h < height; h++ {
|
||
row := data[h*MssWidth*MssBands*PixelBytes : (h+1)*MssWidth*MssBands*PixelBytes]
|
||
for i := 0; i < MssBands; i++ {
|
||
mssData[i] = append(mssData[i], row[i*MssWidth*PixelBytes:(i+1)*MssWidth*PixelBytes]...)
|
||
}
|
||
}
|
||
|
||
for i := 0; i < MssBands; i++ {
|
||
r.MssImages[i], err = gocv.NewMatFromBytes(height, MssWidth, gocv.MatTypeCV16U, mssData[i])
|
||
if err != nil {
|
||
log.Error("Error creating Mat from bytes: ", err)
|
||
return err
|
||
}
|
||
}
|
||
|
||
r.MssHeight = height
|
||
r.MssWidth = MssWidth
|
||
|
||
return nil
|
||
}
|
||
|
||
func (r *Registrator) DoPhaseCorrelation() error {
|
||
// 确保 MSS 高度是 PAN 高度的 1/4
|
||
if r.MssHeight*4 != r.PanHeight {
|
||
err := fmt.Errorf("MSS height is not 1/4 of PAN height, invalid raw file")
|
||
log.Error(err)
|
||
return err
|
||
}
|
||
|
||
// 将PAN将采样作为轮廓匹配基准图像
|
||
downsampledPanImage := gocv.NewMat()
|
||
gocv.Resize(r.PanImage, &downsampledPanImage,
|
||
image.Point{X: r.MssWidth, Y: r.MssHeight}, 0, 0, gocv.InterpolationLinear)
|
||
|
||
// 对每个 MSS 波段图像进行上采样
|
||
upsampledMSSImages := make([]gocv.Mat, MssBands)
|
||
for i := 0; i < MssBands; i++ {
|
||
upsampledMSSImages[i] = gocv.NewMat()
|
||
gocv.Resize(r.MssImages[i], &upsampledMSSImages[i],
|
||
image.Point{X: r.PanWidth, Y: r.PanHeight}, 0, 0, gocv.InterpolationLinear)
|
||
|
||
// r.msToRaw(upsampledMSSImages[i].ToBytes(), fmt.Sprintf("MSS%d.RAW", i+1))
|
||
}
|
||
|
||
fmt.Println("PAN images size:", r.PanImage.Size())
|
||
fmt.Println("upsampled MSS1 images size:", upsampledMSSImages[0].Size())
|
||
fmt.Println("upsampled MSS2 images size:", upsampledMSSImages[0].Size())
|
||
fmt.Println("upsampled MSS3 images size:", upsampledMSSImages[0].Size())
|
||
fmt.Println("upsampled MSS4 images size:", upsampledMSSImages[0].Size())
|
||
|
||
// 分块高度
|
||
blockHeight := r.PanHeight / BlockNum
|
||
|
||
alignedMssData := make([][]byte, MssBands)
|
||
|
||
for band := 0; band < MssBands; band++ {
|
||
var alignedMSSImage gocv.Mat
|
||
for block := 0; block < BlockNum; block++ {
|
||
// 读取 PAN 和 MSS 分块数据
|
||
y1 := (block + 1) * blockHeight
|
||
if block == BlockNum-1 {
|
||
y1 = r.PanHeight
|
||
}
|
||
|
||
log.Println("Processing block", block, image.Rect(0, block*blockHeight, r.PanWidth, y1))
|
||
panBlock := r.PanImage.Region(image.Rect(0, block*blockHeight, r.PanWidth, y1))
|
||
mssBlock := upsampledMSSImages[band].Region(image.Rect(0, block*blockHeight, r.PanWidth, y1))
|
||
|
||
// 处理每个分块
|
||
alignedBlock := r.processBlock(panBlock, mssBlock)
|
||
|
||
alignedBlockData := alignedBlock.ToBytes()
|
||
alignedMssData[band] = append(alignedMssData[band], alignedBlockData...)
|
||
|
||
if alignedMSSImage.Empty() {
|
||
alignedMSSImage = alignedBlock.Clone()
|
||
} else {
|
||
gocv.Vconcat(alignedMSSImage, alignedBlock, &alignedMSSImage)
|
||
alignedBlock.Close()
|
||
}
|
||
|
||
panBlock.Close()
|
||
mssBlock.Close()
|
||
}
|
||
|
||
r.registeredImages[band] = alignedMSSImage
|
||
}
|
||
|
||
r.mssToRaw(alignedMssData, "mss.raw")
|
||
|
||
return nil
|
||
}
|
||
|
||
func (r *Registrator) processBlock(panBlock, mssBlock gocv.Mat) gocv.Mat {
|
||
pan := gocv.NewMat()
|
||
mss := gocv.NewMat()
|
||
|
||
panBlock.ConvertTo(&pan, gocv.MatTypeCV32FC1)
|
||
defer pan.Close()
|
||
mssBlock.ConvertTo(&mss, gocv.MatTypeCV32FC1)
|
||
defer mss.Close()
|
||
|
||
hann := gocv.NewMat()
|
||
shift, response := gocv.PhaseCorrelate(pan, mss, hann)
|
||
hann.Close()
|
||
|
||
dx := shift.X
|
||
dy := shift.Y
|
||
log.Printf("Block shift: dx = %f, dy = %f. response = %f \n", dx, dy, response)
|
||
|
||
// cv::Mat translationMatrix = (cv::Mat_<double>(2, 3) << 1, 0, shift.x, 0, 1, shift.y);
|
||
translationMat := gocv.NewMatWithSize(2, 3, gocv.MatTypeCV32F)
|
||
translationMat.SetFloatAt(0, 0, 1)
|
||
translationMat.SetFloatAt(0, 1, 0)
|
||
translationMat.SetFloatAt(0, 2, shift.X)
|
||
translationMat.SetFloatAt(1, 0, 0)
|
||
translationMat.SetFloatAt(1, 1, 1)
|
||
translationMat.SetFloatAt(1, 2, shift.Y)
|
||
|
||
alignedBlock := gocv.NewMat()
|
||
gocv.WarpAffine(mss, &alignedBlock, translationMat, image.Pt(panBlock.Size()[0], panBlock.Size()[1]))
|
||
|
||
translationMat.Close()
|
||
alignedBlock.ConvertTo(&alignedBlock, gocv.MatTypeCV16U)
|
||
|
||
return alignedBlock
|
||
}
|
||
|
||
func (r *Registrator) panToTiff(panImage gocv.Mat, filePath string) error {
|
||
return nil
|
||
}
|
||
|
||
func (r *Registrator) mssToTiff(registeredImages [4]gocv.Mat, filePath string) error {
|
||
// 创建合并后的图像（RGBIR）
|
||
rgbirImage := gocv.NewMatWithSize(r.PanHeight, r.PanWidth, gocv.MatTypeCV16UC4) // 4通道，16位
|
||
|
||
for y := 0; y < r.PanHeight; y++ {
|
||
for x := 0; x < r.PanWidth; x++ {
|
||
r := registeredImages[0].GetShortAt(y, x)
|
||
g := registeredImages[1].GetShortAt(y, x)
|
||
b := registeredImages[2].GetShortAt(y, x)
|
||
ir := registeredImages[3].GetShortAt(y, x)
|
||
rgbirImage.SetShortAt(y, x*4+0, r)
|
||
rgbirImage.SetShortAt(y, x*4+1, g)
|
||
rgbirImage.SetShortAt(y, x*4+2, b)
|
||
rgbirImage.SetShortAt(y, x*4+3, ir)
|
||
}
|
||
}
|
||
|
||
// 将合并后的图像保存为TIFF文件
|
||
fileName := "registered_rgbir.tiff"
|
||
tiffFile, err := os.Create(fileName)
|
||
if err != nil {
|
||
fmt.Println("Error creating TIFF file:", err)
|
||
return err
|
||
}
|
||
defer tiffFile.Close()
|
||
|
||
// 使用tiff库保存图像
|
||
img, err := rgbirImage.ToImage()
|
||
if err != nil {
|
||
fmt.Println("Error converting Mat to image:", err)
|
||
return err
|
||
}
|
||
if err := tiff.Encode(tiffFile, img, nil); err != nil {
|
||
fmt.Println("Error encoding TIFF file:", err)
|
||
return err
|
||
}
|
||
|
||
fmt.Println("Saved", fileName)
|
||
return nil
|
||
}
|
||
|
||
func (r *Registrator) mssToRaw(mssData [][]byte, filePath string) error {
|
||
f, err := os.OpenFile("mss.RAW", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0777)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
w := bufio.NewWriter(f)
|
||
for row := 0; row < r.PanHeight; row++ {
|
||
w.Write(mssData[0][row*r.PanWidth*PixelBytes : (row+1)*r.PanWidth*PixelBytes])
|
||
w.Write(mssData[1][row*r.PanWidth*PixelBytes : (row+1)*r.PanWidth*PixelBytes])
|
||
w.Write(mssData[2][row*r.PanWidth*PixelBytes : (row+1)*r.PanWidth*PixelBytes])
|
||
w.Write(mssData[3][row*r.PanWidth*PixelBytes : (row+1)*r.PanWidth*PixelBytes])
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
func (r *Registrator) msToRaw(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
|
||
}
|