Files
sjy01-image-proc/pkg/fusion/fusion.go
2024-11-11 17:16:50 +08:00

269 lines
6.5 KiB
Go

package fusion
import (
"image"
"os"
"strings"
"github.com/airbusgeo/godal"
log "github.com/sirupsen/logrus"
"gocv.io/x/gocv"
"starwiz.cn/sjy01/image-proc/pkg/config"
"starwiz.cn/sjy01/image-proc/pkg/utils"
)
// Reference: https://www.kaggle.com/code/resolut/panchromatic-sharpening
type PansharpenMethod int
const (
SimpleBrovey PansharpenMethod = iota
Brovey
SimpleMean
ESRI
GDALPansharpen
IHS
GIHS
WeightR = 1
WeightG = 0.75
WeightB = 0.55
WeightNIR = 0.85
)
func Pansharpen(panTif, mssTif, fusTif string, method PansharpenMethod, weight float32) error {
log.Infof("Pansharpening %s and %s", panTif, mssTif)
pan, err := utils.ReadTifftoMat(panTif)
if err != nil {
return err
}
defer pan.Close()
log.Infof("PAN row %d, col %d", pan.Rows(), pan.Cols())
mss, err := utils.ReadTifftoMat(mssTif)
if err != nil {
return err
}
defer mss.Close()
log.Infof("MSS bands %d row %d, col %d", mss.Channels(), mss.Rows(), mss.Cols())
mssResized := gocv.NewMat()
defer mssResized.Close()
gocv.Resize(*mss, &mssResized, image.Point{X: pan.Cols(), Y: pan.Rows()}, 0, 0, gocv.InterpolationCubic)
BGRI := gocv.Split(mssResized)
var fus gocv.Mat
switch method {
case ESRI:
fus = PansharpenESRI(*pan, BGRI[0], BGRI[1], BGRI[2], BGRI[3])
case IHS:
fus = PansharpenIHS(*pan, BGRI[0], BGRI[1], BGRI[2], BGRI[3], weight)
case Brovey:
fus = PansharpenBrovey(*pan, BGRI[0], BGRI[1], BGRI[2], BGRI[3], weight)
default:
fus = PansharpenESRI(*pan, BGRI[0], BGRI[1], BGRI[2], BGRI[3])
}
defer fus.Close()
var mt gocv.MatType
switch fus.Channels() {
case 1:
mt = gocv.MatTypeCV16UC1
case 3:
mt = gocv.MatTypeCV16UC3
case 4:
mt = gocv.MatTypeCV16UC4
}
fus.ConvertTo(&fus, mt)
var colorInterps []godal.ColorInterp
if config.GCONFIG.CoRegistration.FUSBandOrder == "RGB" ||
config.GCONFIG.CoRegistration.FUSBandOrder == "RGBI" {
colorInterps = []godal.ColorInterp{godal.CIRed, godal.CIGreen, godal.CIBlue, godal.CIUndefined}
} else {
colorInterps = []godal.ColorInterp{godal.CIBlue, godal.CIGreen, godal.CIRed, godal.CIUndefined}
}
err = utils.SaveBGRToGDALGTiff(fus, fus.Channels(), 0, 0, 1.3, colorInterps, fusTif)
for i := 0; i < len(BGRI); i++ {
BGRI[i].Close()
}
// use rpc of PAN directly
rpb := strings.Replace(panTif, ".tiff", ".rpb", -1)
if _, err := os.Stat(rpb); err == nil {
fusRPB := strings.Replace(fusTif, ".tiff", ".rpb", -1)
if input, err := os.ReadFile(rpb); err == nil {
os.WriteFile(fusRPB, input, 0644)
}
}
return err
}
func PansharpenSimpleBrovey(pan, B, G, R, I gocv.Mat) error {
return nil
}
func PansharpenBrovey(pan, B, G, R, I gocv.Mat, w float32) gocv.Mat {
pan.ConvertTo(&pan, gocv.MatTypeCV32F)
B.ConvertTo(&B, gocv.MatTypeCV32F)
G.ConvertTo(&G, gocv.MatTypeCV32F)
R.ConvertTo(&R, gocv.MatTypeCV32F)
I.ConvertTo(&I, gocv.MatTypeCV32F)
// DNF = (PAN - w*I)/(R*w + G*w + B*w)
wI := gocv.NewMat()
I.CopyTo(&wI)
wI.MultiplyFloat(-1 * w)
defer wI.Close()
wR := gocv.NewMat()
R.CopyTo(&wR)
wR.MultiplyFloat(w)
defer wR.Close()
wG := gocv.NewMat()
G.CopyTo(&wG)
wG.MultiplyFloat(w)
defer wG.Close()
wB := gocv.NewMat()
B.CopyTo(&wB)
wB.MultiplyFloat(w)
defer wB.Close()
DNF := gocv.NewMat()
defer DNF.Close()
gocv.Add(pan, wI, &DNF)
gocv.Add(wR, wG, &wR)
gocv.Add(wR, wB, &wR)
gocv.Divide(DNF, wR, &DNF)
gocv.Multiply(R, DNF, &R)
gocv.Multiply(G, DNF, &G)
gocv.Multiply(B, DNF, &B)
gocv.Multiply(I, DNF, &I)
mss := gocv.NewMat()
mt := gocv.MatTypeCV16UC3
switch config.GCONFIG.CoRegistration.FUSBandOrder {
case "RGB":
gocv.Merge([]gocv.Mat{R, G, B}, &mss)
case "BGR":
gocv.Merge([]gocv.Mat{B, G, R}, &mss)
case "RGBI":
gocv.Merge([]gocv.Mat{R, G, B, I}, &mss)
mt = gocv.MatTypeCV16UC4
case "BGRI":
gocv.Merge([]gocv.Mat{B, G, R, I}, &mss)
mt = gocv.MatTypeCV16UC4
default:
gocv.Merge([]gocv.Mat{R, G, B}, &mss)
}
mss.ConvertTo(&mss, mt)
return mss
}
func PansharpenSimpleMean(pan, B, G, R, I gocv.Mat) error { return nil }
func PansharpenESRI(pan, B, G, R, I gocv.Mat) gocv.Mat {
bgrn_mean := gocv.NewMatWithSize(pan.Rows(), pan.Cols(), gocv.MatTypeCV32F)
defer bgrn_mean.Close()
ADJ := gocv.NewMat()
defer ADJ.Close()
pan.ConvertTo(&pan, gocv.MatTypeCV32F)
B.ConvertTo(&B, gocv.MatTypeCV32F)
G.ConvertTo(&G, gocv.MatTypeCV32F)
R.ConvertTo(&R, gocv.MatTypeCV32F)
I.ConvertTo(&I, gocv.MatTypeCV32F)
for y := 0; y < pan.Rows(); y++ {
for x := 0; x < pan.Cols(); x++ {
v := (R.GetFloatAt(y, x) + G.GetFloatAt(y, x) + B.GetFloatAt(y, x) + I.GetFloatAt(y, x)) / 4
bgrn_mean.SetFloatAt(y, x, v)
}
}
bgrn_mean.MultiplyFloat(-1)
gocv.Add(pan, bgrn_mean, &ADJ)
gocv.Add(B, ADJ, &B)
gocv.Add(G, ADJ, &G)
gocv.Add(R, ADJ, &R)
gocv.Add(I, ADJ, &I)
mss := gocv.NewMat()
mt := gocv.MatTypeCV16UC3
switch config.GCONFIG.CoRegistration.FUSBandOrder {
case "RGB":
gocv.Merge([]gocv.Mat{R, G, B}, &mss)
case "BGR":
gocv.Merge([]gocv.Mat{B, G, R}, &mss)
case "RGBI":
gocv.Merge([]gocv.Mat{R, G, B, I}, &mss)
mt = gocv.MatTypeCV16UC4
case "BGRI":
gocv.Merge([]gocv.Mat{B, G, R, I}, &mss)
mt = gocv.MatTypeCV16UC4
default:
gocv.Merge([]gocv.Mat{R, G, B}, &mss)
}
mss.ConvertTo(&mss, mt)
return mss
}
func PansharpenIHS(pan, B, G, R, I gocv.Mat, w float32) gocv.Mat {
mss := gocv.NewMat()
gocv.Merge([]gocv.Mat{B, G, R}, &mss)
mssHLS := gocv.NewMat()
defer mssHLS.Close()
mss.ConvertTo(&mss, gocv.MatTypeCV32FC3)
gocv.CvtColor(mss, &mssHLS, gocv.ColorBGRToHLS)
hlsChannels := gocv.Split(mssHLS)
// 将 PAN 图像转换为与 HLS 亮度分量相同的范围
panFloat := gocv.NewMat()
defer panFloat.Close()
pan.ConvertTo(&panFloat, gocv.MatTypeCV32F)
gocv.Normalize(panFloat, &panFloat, 0, 65535, gocv.NormMinMax)
hlsFloat := gocv.NewMat()
defer hlsFloat.Close()
hlsChannels[1].ConvertTo(&hlsFloat, gocv.MatTypeCV32F)
gocv.Normalize(hlsFloat, &hlsFloat, 0, 65535, gocv.NormMinMax)
// 替换亮度分量
I.ConvertTo(&I, gocv.MatTypeCV32F)
I.MultiplyFloat(-1 * w)
gocv.Add(panFloat, I, &hlsChannels[1])
// 合并通道
gocv.Merge(hlsChannels, &mssHLS)
// 将图像从 HLS 转换回 BGR
fusedImage := gocv.NewMat()
gocv.CvtColor(mssHLS, &fusedImage, gocv.ColorHLSToBGR)
channels := gocv.Split(fusedImage)
defer channels[0].Close() // B
defer channels[1].Close() // G
defer channels[2].Close() // R
if config.GCONFIG.CoRegistration.FUSBandOrder == "RGBI" || config.GCONFIG.CoRegistration.FUSBandOrder == "RGB" {
gocv.Merge([]gocv.Mat{channels[2], channels[1], channels[0]}, &fusedImage)
} else {
gocv.Merge([]gocv.Mat{channels[0], channels[1], channels[2]}, &fusedImage)
}
fusedImage.ConvertTo(&fusedImage, gocv.MatTypeCV16UC3)
return fusedImage
}