sjy01-image-proc/output.go

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.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
}

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.MssWidth * PixelBytes
	height := r.MssHeight
	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: r.MssWidth,
		Lines:   r.MssHeight,
		Bands:   MssBands,
	}

	hdrFile := filepath.Join(filepath.Dir(raw),
		fmt.Sprintf("%s.HDR", strings.TrimSuffix(filepath.Base(raw), filepath.Ext(raw))),
	)
	hdr.Write(hdrFile)

	return nil
}