使用gamma校正提升jpg亮度

2024-05-30 11:22:34 +08:00
parent 7d9ec46750
commit 07ee4d88d4
18 changed files with 174 additions and 98 deletions
--- a/pkg/producer/output.go
+++ b/pkg/producer/output.go
@@ -0,0 +1,196 @@
+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 {
+	err := savePanToGDALGTiff(r.PanImage, tiffFile)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+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 {
+	r.rgbirImage = gocv.NewMat()
+	gocv.Merge(r.registeredMssImages[:], &r.rgbirImage)
+
+	err := SaveBGRToGDALGTiff(r.rgbirImage,
+		4, 5,
+		[]godal.ColorInterp{godal.CIBlue, godal.CIGreen, godal.CIRed, godal.CIUndefined},
+		tiffFile)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+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 {
+	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
+}