sjy01-image-proc/cmd/proc.go

package main

import (
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/airbusgeo/godal"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"starwiz.cn/sjy01/image-proc/pkg/config"
	producer "starwiz.cn/sjy01/image-proc/pkg/producer"
)

var (
	params           producer.Params
	saveStrip        bool
	doLUTRRC         bool
	doMomentMatching bool
	lutDir           string
)

var procCmd = &cobra.Command{
	Use:   "proc",
	Short: "process images",
	Long:  `process images`,
	Run: func(cmd *cobra.Command, args []string) {
		config.GViper = config.InitViper(configFile)
		logrus.SetLevel(config.GCONFIG.LogLevel)

		reg := producer.NewRegistrator(producer.DownSampled)
		reg.Params = initParams()

		if err := reg.LoadAuxData(); err != nil {
			logrus.Fatal(err)
		}
		reg.AuxPrint()

		if err := reg.LoadMssRaw(); err != nil {
			logrus.Fatal(err)
		}

		if err := reg.LoadPanRaw(); err != nil {
			logrus.Fatal(err)
		}

		godal.RegisterAll()
		os.MkdirAll(params.OutputDir, 0755)

		if doLUTRRC {
			reg.DoRRCbyLUT(lutDir)
		}

		if doMomentMatching {
			reg.DoMomentMatching()
		}

		if err := reg.DoPhaseCorrelation(); err != nil {
			logrus.Fatal(err)
		}
		reg.DoCoRegistration()

		if params.SaveRegisteredMssRaw {
			registerdMSSRAW := filepath.Join(
				params.OutputDir,
				strings.TrimSuffix(filepath.Base(params.MssRawFile), filepath.Ext(params.MssRawFile))+"_registered.RAW",
			)
			reg.SaveRegisteredMssToRaw(registerdMSSRAW)
		}

		panScenes, mssScenes, err := reg.SubScenes()
		if err != nil {
			logrus.Error(err)
		}

		reg.SaveScenesToTiff(panScenes, mssScenes)
		producer.CleanScenes(panScenes)
		producer.CleanScenes(mssScenes)
		runtime.GC()

		if saveStrip {
			reg.SaveOriginalPanToGDALGTiff(reg.Params.PanTiffFile)
			reg.SaveRegisteredMssToGDALGTiff(reg.Params.MssTiffFile)
		}

		if reg.Params.DoPansharpen {
			reg.DoScenePansharpen(panScenes, mssScenes)
		}

		reg.Report()

		reg.Clean()
	},
}

func init() {
	rootCmd.AddCommand(procCmd)

	procCmd.Flags().StringVarP(&params.PanRawFile, "pan", "p", "", "PAN image raw file path")
	procCmd.Flags().StringVarP(&params.MssRawFile, "mss", "m", "", "MSS image raw file path")
	procCmd.Flags().StringVarP(&params.AuxRawFile, "aux", "a", "", "AUX image raw file path")
	procCmd.Flags().BoolVarP(&params.SaveRegisteredMssRaw, "srmss", "s", false, "save registered MSS image raw file")
	procCmd.Flags().BoolVarP(&params.DoPansharpen, "fus", "", false, "pansharpen using IHS")
	procCmd.Flags().StringVarP(&params.OutputDir, "output-dir", "o", "data", "output directory")
	procCmd.Flags().BoolVarP(&params.SubScenes, "sub-scenes", "", false, "process sub-scenes")
	procCmd.Flags().BoolVarP(&saveStrip, "save-strip", "", false, "save original and registered images as GDAL GTiff")
	procCmd.Flags().StringVarP(&paramsXML, "params", "x", "", "params xml file path")
	procCmd.Flags().StringVarP(&lutDir, "lut-dir", "", "data/lut", "LUT directory")
	procCmd.Flags().BoolVarP(&doLUTRRC, "lut-rrc", "", false, "do RRC with LUT method")
	procCmd.Flags().BoolVarP(&doMomentMatching, "mm", "", false, "do moment matching")
}

func initParams() producer.Params {
	taskParams := params

	if paramsXML != "" {

		task, err := producer.ParseXMLImageTask(paramsXML)
		if err != nil {
			logrus.Fatal(err)
		}

		taskParams.PanRawFile = task.InputFileList.PanData
		taskParams.MssRawFile = task.InputFileList.MssData
		taskParams.AuxRawFile = task.InputFileList.AuxData
		taskParams.DoPansharpen = task.Params.DoPansharpen
		taskParams.OutputDir = task.Params.OutputPath
		taskParams.ReportFile = task.Params.ReportFile
		taskParams.DataId = task.Params.DataID
	}

	taskParams.MssTiffFile = filepath.Join(taskParams.OutputDir, strings.TrimSuffix(filepath.Base(taskParams.MssRawFile),
		filepath.Ext(taskParams.MssRawFile))+".tiff")

	taskParams.PanTiffFile = filepath.Join(taskParams.OutputDir,
		strings.TrimSuffix(filepath.Base(taskParams.PanRawFile),
			filepath.Ext(taskParams.PanRawFile))+".tiff")

	if taskParams.ReportFile == "" {
		id := strings.TrimSuffix(filepath.Base(taskParams.PanRawFile), filepath.Ext(taskParams.PanRawFile))
		id = strings.Replace(id, "PAN", "PMS", 1)
		taskParams.ReportFile = filepath.Join(taskParams.OutputDir, id+"_report.xml")
	}

	return taskParams
}