sjy01-image-proc/pkg/utils/tiff.go

package utils

import (
	"os"
	"path/filepath"
	"time"

	"github.com/airbusgeo/godal"
	log "github.com/sirupsen/logrus"
	"starwiz.cn/sjy01/image-proc/pkg/calculator"

	"gocv.io/x/gocv"
)

func SavePanToGDALGTiff(pan gocv.Mat, topLeftX, topLeftY float64, tiffFile string, resolution float64) error {
	// log.Println("Saving PAN image to TIFF file:", tiffFile)
	os.MkdirAll(filepath.Dir(tiffFile), 0755)
	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, topLeftX, topLeftY, resolution)
	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 image to ", tiffFile)

	return nil
}

func SaveBGRToGDALGTiff(bgr gocv.Mat,
	bands int,
	topLeftX, topLeftY float64,
	resolution float64,
	colorInterps []godal.ColorInterp, tiffFile string) error {
	width := bgr.Cols()
	height := bgr.Rows()

	os.MkdirAll(filepath.Dir(tiffFile), 0755)

	// 创建一个二维切片来存储图像数据
	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, topLeftX, topLeftY, 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 setGeoTransform(ds *godal.Dataset, topLeftLng, topLeftLat, resolution float64) {
	// 转换分辨率（米到度）
	resLat := resolution / calculator.MetersPerDegreeLatitude
	resLng := calculator.MetersToDegreesLongitude(resolution, topLeftLat)
	// 设置地理变换（假设左上角坐标为(0,0)，PAN每个像素分辨率为1.2米）
	geotransform := [6]float64{
		topLeftLng, // top left x
		resLng,     // w-e pixel resolution
		0,          // rotation, 0 if image is "north up"
		topLeftLat, // top left y
		0,          // rotation, 0 if image is "north up"
		-resLat,    // n-s pixel resolution (negative value)
	}
	err := ds.SetGeoTransform(geotransform)
	if err != nil {
		log.Errorf("Failed to set GeoTransform: %v", err)
		return
	}

	// 设置投影为 WGS84
	srs, err := godal.NewSpatialRefFromEPSG(4326) //  gdal.CreateSpatialReference("")
	if err != nil {
		log.Errorf("Failed to set spatial reference: %v", err)
		return
	}

	projWKT, err := srs.WKT()
	if err != nil {
		log.Errorf("Failed to convert spatial reference to WKT: %v", err)
		return
	}

	err = ds.SetProjection(projWKT)
	if err != nil {
		log.Errorf("Failed to set projection: %v", err)
	}

	// 设置一些常见的元数据（可选）
	ds.SetMetadata("TIFFTAG_DATETIME", time.Now().String())
	ds.SetMetadata("TIFFTAG_SOFTWARE", "StarWiz-SJY01-IMAGE-PROC")
}

func ReadTiff(tifFile string) ([][][]float32, error) {
	hDataset, err := godal.Open(tifFile)
	if err != nil {
		return nil, err
	}

	var data [][][]float32
	for _, band := range hDataset.Bands() {
		structure := band.Structure()
		var dataBand [][]float32
		for y := 1; y <= structure.SizeY; y++ {
			scanline := make([]float32, structure.SizeX)
			err = band.Read(0, y, scanline, structure.SizeX, 1)
			if err != nil {
				return nil, err
			}
			dataBand = append(dataBand, scanline)
		}

		data = append(data, dataBand)
	}
	return data, nil
}

func ReadTifftoMat(ftiff string) (*gocv.Mat, error) {
	ds, err := godal.Open(ftiff)
	if err != nil {
		log.Printf("Error opening TIFF file %s: %v", ftiff, err)
		return nil, err
	}
	defer ds.Close()

	bands := ds.Bands()

	// 获取图像大小
	width := bands[0].Structure().SizeX
	height := bands[0].Structure().SizeY
	bandsCnt := len(bands)

	var img gocv.Mat
	if bandsCnt == 1 {
		img = gocv.NewMatWithSize(height, width, gocv.MatTypeCV16UC1)
	} else {
		img = gocv.NewMatWithSize(height, width, gocv.MatTypeCV16UC4)
	}

	// 读取每个波段并存储到图像矩阵中
	for i := 0; i < bandsCnt; i++ {
		band := bands[i]
		data := make([]uint16, width*height)
		err = band.Read(0, 0, data, width, height)
		if err != nil {
			return nil, err
		}

		// 将 16 位数据存储到图像的相应通道
		for y := 0; y < height; y++ {
			for x := 0; x < width; x++ {
				value := data[y*width+x]
				img.SetShortAt(y, x*bandsCnt+i, int16(value))
			}
		}
	}

	return &img, nil
}