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/calculator"
	"starwiz.cn/sjy01/image-proc/pkg/config"
	"starwiz.cn/sjy01/image-proc/pkg/dem"
	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)
		os.MkdirAll(config.GCONFIG.Log.LogDir, 0755)
		configureLogger(logrus.StandardLogger(),
			filepath.Join(config.GCONFIG.Log.LogDir, "imgproc.log"),
			config.GCONFIG.Log.LogLevel)

		logrus.SetLevel(config.GCONFIG.Log.LogLevel)
		logrus.Info("image-proc version:", Version)

		godal.RegisterAll()

		// dem
		dem.Dem1KmLT = dem.NewDem1Km(config.GCONFIG.Dem.Dem1Km)

		calculator.EOP = calculator.NewEOPTable()
		calculator.EOP.Load(eopData, eopp5Line)

		processor := producer.NewImgProc(producer.DownSampled)
		processor.Params = initParams()

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

		// reg.AuxPrint()

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

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

		// for i, img := range reg.MssImages {
		// 	edge := producer.CV_Sobel(img)
		// 	utils.SavePanToGDALGTiff(edge, 0, 0, "data/temp/out_"+strconv.Itoa(i)+".tif", 1.3)
		// }
		// return

		os.MkdirAll(params.OutputDir, 0755)

		if doLUTRRC {
			processor.DoRRCbyLUT(lutDir)
		}

		if doMomentMatching {
			processor.DoMomentMatching()
		}

		if err := processor.DoPhaseCorrelation(false); err != nil {
			logrus.Fatal(err)
		}

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

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

		// scenesCnt := mathutil.Min(len(panScenes), len(mssScenes))
		// for i := 0; i < scenesCnt; i++ {
		// 	processor.RegisterScenes(panScenes[i], mssScenes[i])
		// }

		processor.OutputL1A(panScenes, mssScenes)
		producer.CleanScenes(panScenes)
		producer.CleanScenes(mssScenes)

		runtime.GC()

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

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

		processor.Report()

		processor.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.Targets = task.Params.Targets
	}

	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
}