sjy01-image-proc/pkg/producer/scenes.go

package imageproc

import (
	"fmt"
	"image"
	"os"
	"path/filepath"
	"strings"

	"github.com/airbusgeo/godal"
	log "github.com/sirupsen/logrus"
	"gocv.io/x/gocv"
)

type Scene struct {
	Width  int
	Height int
	X      int // coordinate of the left top corner of the scene
	Y      int
	Mat    []gocv.Mat
	Tiff   string
}

func (s *Scene) Cleanup() {
	for i := range s.Mat {
		s.Mat[i].Close()
	}
}

// 对 PAN 和 配准后的MSS 在 Y 方向进行分景，景之间有25%的重叠
// 默认分景大小：
// MSS  2336 * 2336
// PAN  9344 * 9344
func (r *Registrator) SubScenes() (panScenes []*Scene, mssScenes []*Scene, err error) {
	hPAN := (9344 - 2336)
	hPANOverlap := 2336
	panScenesCnt := r.PanHeight / hPAN

	for i := 0; i < panScenesCnt; i++ {
		y1 := (i+1)*hPAN + hPANOverlap
		if y1 > r.PanHeight {
			y1 = r.PanHeight
		}

		scene := &Scene{
			Width:  9344,
			Height: y1 - i*hPAN,
			X:      0,
			Y:      i * hPAN,
		}

		mat := r.PanImage.Region(image.Rect(0, i*hPAN, 9344, y1))
		scene.Mat = append(scene.Mat, mat)
		panScenes = append(panScenes, scene)
	}

	hMSS := hPAN / 4
	hMSSOverlap := hPANOverlap / 4
	mssScenesCnt := r.MssHeight / hMSS
	for i := 0; i < mssScenesCnt; i++ {
		y1 := (i+1)*hMSS + hMSSOverlap
		if y1 > r.MssHeight {
			y1 = r.MssHeight
		}

		scene := &Scene{
			Width:  2336,
			Height: y1 - i*hMSS,
			X:      0,
			Y:      i * hMSS,
		}

		for band := 0; band < 4; band++ {
			mat := r.registeredMssImages[band].Region(image.Rect(0, i*hMSS, 2336, y1))
			scene.Mat = append(scene.Mat, mat)
		}

		mssScenes = append(mssScenes, scene)
	}

	if len(panScenes) != len(mssScenes) {
		log.Error("pan and mss scenes count not match")
		err = fmt.Errorf("pan and mss scenes count not match")
	}

	return panScenes, mssScenes, err
}

func (r *Registrator) SaveScenesToTiff(panScenes []*Scene, mssScenes []*Scene) error {
	for i, scene := range panScenes {
		dir := filepath.Join(r.Params.OutputDir, fmt.Sprintf("%03d", i+1), "PAN")
		os.MkdirAll(dir, 0755)

		tiff := filepath.Base(r.Params.PanTiffFile)
		tiff = strings.TrimSuffix(tiff, ".tiff")
		filename := fmt.Sprintf("%s_%03d.tiff", tiff, i+1)
		scene.Tiff = filepath.Join(dir, filename)
		err := savePanToGDALGTiff(scene.Mat[0], scene.Tiff)
		if err != nil {
			log.Errorf("save pan scene %d to tiff failed: %v", i+1, err)
			return err
		}

		GTiffToJPG(scene.Tiff, strings.Replace(scene.Tiff, ".tiff", ".jpg", 1), false)
	}

	for i, scene := range mssScenes {
		dir := filepath.Join(r.Params.OutputDir, fmt.Sprintf("%03d", i+1), "MSS")
		os.MkdirAll(dir, 0755)

		tiff := filepath.Base(r.Params.MssTiffFile)
		tiff = strings.TrimSuffix(tiff, ".tiff")
		filename := fmt.Sprintf("%s_%03d.tiff", tiff, i+1)
		scene.Tiff = filepath.Join(dir, filename)

		rgbirImage, _ := r.MergeMSSToBGRNIR(scene.Mat)
		err := SaveBGRToGDALGTiff(rgbirImage,
			4, 5,
			[]godal.ColorInterp{godal.CIBlue, godal.CIGreen, godal.CIRed, godal.CIUndefined},
			scene.Tiff)
		if err != nil {
			log.Errorf("save mss scene %d to tiff failed: %v", i+1, err)
			return err
		}

		GTiffToJPG(scene.Tiff, strings.Replace(scene.Tiff, ".tiff", ".jpg", 1), false)
	}

	return nil
}

func (r *Registrator) DoScenePansharpen(panScenes []*Scene, mssScenes []*Scene) error {
	for i := 0; i < len(panScenes); i++ {
		// rgbirImg, _ := r.MergeMSSToBGRNIR(mssScenes[i].Mat)
		// ihsImage := PansharpenIHS(panScenes[i].Mat[0], rgbirImg)
		// path := strings.TrimSuffix(r.Params.MssTiffFile, ".tiff")
		// filename := fmt.Sprintf("%s_%03d_FUS.tiff", path, i+1)

		// SaveBGRToGDALGTiff(ihsImage, 3, 1.25, filename)
		fusedTiff, err := GDALPansharpen(panScenes[i].Tiff, mssScenes[i].Tiff)
		if err != nil {
			return err
		}

		GTiffToJPG(fusedTiff, strings.Replace(fusedTiff, ".tiff", ".jpg", 1), true)
	}
	return nil
}

func (r *Registrator) MergeMSSToBGRNIR(channels []gocv.Mat) (gocv.Mat, error) {
	var rgbirImage gocv.Mat
	if len(channels) != 4 {
		return rgbirImage, fmt.Errorf("mss channels count not match")
	}

	rgbirImage = gocv.NewMat()
	gocv.Merge(channels, &rgbirImage)
	log.Printf("merge mss to bgr nir image, channels: %d, height: %d, width: %d",
		rgbirImage.Channels(), rgbirImage.Rows(), rgbirImage.Cols())

	return rgbirImage, nil

}