parametres input

2024-05-28 00:03:18 +08:00
parent 720a3fd855
commit c7bbb632f5
10 changed files with 245 additions and 90 deletions
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -1,9 +1,7 @@
 package main

 import (
-	"github.com/airbusgeo/godal"
 	"github.com/sirupsen/logrus"
-	imageproc "starwiz.cn/sjy01/image-proc"
 )

 func init() {
@@ -16,29 +14,5 @@ func init() {
 }

 func main() {
-	reg := imageproc.NewRegistrator()
-
-	if err := reg.LoadMssRaw("/Users/lan/workspace/temp/ccv/ms_image.raw"); err != nil {
-		panic(err)
-	}
-
-	if err := reg.LoadPanRaw("/Users/lan/workspace/temp/ccv/pan_image.raw"); err != nil {
-		panic(err)
-	}
-
-	godal.RegisterAll()
-
-	// reg.SaveOriginalPanToGDALGTiff("data/original_pan.tiff")
-	// reg.SaveFilteredPanToGDALGTiff("data/filtered_pan.tiff")
-
-	if err := reg.DoPhaseCorrelation(); err != nil {
-		panic(err)
-	}
-	reg.DoCoRegestration()
-
-	// // reg.SaveRegisteredMssToRaw("data/registered_mss.RAW")
-	reg.SaveRegisteredMssToGDALGTiff("data/registered_mss.tiff")
-	reg.SavePansharpenedToGDALGTiff("data/pansharpened.tiff")
-
-	reg.Clean()
+	rootCmd.Execute()
 }
--- a/cmd/proc.go
+++ b/cmd/proc.go
@@ -0,0 +1,88 @@
+package main
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/airbusgeo/godal"
+	"github.com/spf13/cobra"
+	imageproc "starwiz.cn/sjy01/image-proc"
+)
+
+var (
+	PanRawFile           string
+	MssRawFile           string
+	AuxRawFile           string
+	SaveRegisteredMssRaw bool
+	PansharpenIHS        bool
+	OutputDir            string
+)
+
+var procCmd = &cobra.Command{
+	Use:   "proc",
+	Short: "process images",
+	Long:  `process images`,
+	Run: func(cmd *cobra.Command, args []string) {
+		reg := imageproc.NewRegistrator(imageproc.DownSampled)
+		reg.Params = imageproc.Params{
+			PanRawFile:           PanRawFile,
+			MssRawFile:           MssRawFile,
+			AuxRawFile:           AuxRawFile,
+			SaveRegisteredMssRaw: SaveRegisteredMssRaw,
+			PansharpenIHS:        PansharpenIHS,
+			OutputDir:            OutputDir,
+			MssTiffFile: filepath.Join(OutputDir,
+				strings.TrimSuffix(filepath.Base(MssRawFile),
+					filepath.Ext(MssRawFile))+".tiff"),
+			PanTiffFile: filepath.Join(OutputDir,
+				strings.TrimSuffix(filepath.Base(PanRawFile),
+					filepath.Ext(PanRawFile))+".tiff"),
+		}
+		reg.Params.FusTIffFile = strings.Replace(reg.Params.MssTiffFile, ".tiff", "_FUS.tiff", 1)
+
+		if err := reg.LoadMssRaw(); err != nil {
+			panic(err)
+		}
+
+		if err := reg.LoadPanRaw(); err != nil {
+			panic(err)
+		}
+
+		godal.RegisterAll()
+		os.MkdirAll(OutputDir, 0755)
+
+		if err := reg.DoPhaseCorrelation(); err != nil {
+			panic(err)
+		}
+		reg.DoCoRegestration()
+
+		if SaveRegisteredMssRaw {
+			registerdMSSRAW := filepath.Join(
+				OutputDir,
+				strings.TrimSuffix(filepath.Base(MssRawFile), filepath.Ext(MssRawFile))+"_registered.RAW",
+			)
+			reg.SaveRegisteredMssToRaw(registerdMSSRAW)
+		}
+
+		reg.SaveOriginalPanToGDALGTiff(reg.Params.PanTiffFile)
+		reg.SaveRegisteredMssToGDALGTiff(reg.Params.MssTiffFile)
+
+		if PansharpenIHS {
+			reg.SavePansharpenedToGDALGTiff(reg.Params.FusTIffFile)
+		}
+
+		reg.Clean()
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(procCmd)
+
+	procCmd.Flags().StringVarP(&PanRawFile, "pan", "p", "pan.raw", "PAN image raw file path")
+	procCmd.Flags().StringVarP(&MssRawFile, "mss", "m", "mss.raw", "MSS image raw file path")
+	procCmd.Flags().StringVarP(&AuxRawFile, "aux", "a", "pms.aux", "AUX image raw file path")
+	procCmd.Flags().BoolVarP(&SaveRegisteredMssRaw, "srmss", "s", false, "save registered MSS image raw file")
+	procCmd.Flags().BoolVarP(&PansharpenIHS, "fus", "", false, "pansharpen using IHS")
+	procCmd.Flags().StringVarP(&OutputDir, "output-dir", "o", "data", "output directory")
+}
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -0,0 +1,16 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+)
+
+var rootCmd = &cobra.Command{
+	Use:   "SJY01 Image processing tools",
+	Short: "Preprocessing tools for SJY01 original image data",
+	Long:  ``,
+	Run: func(cmd *cobra.Command, args []string) {
+		fmt.Println("preprocessing tools for SJY01 original image data")
+	},
+}
--- a/fit.go
+++ b/fit.go
@@ -1,13 +1,11 @@
 package imageproc

 import (
-	"fmt"
-
 	"gonum.org/v1/gonum/mat"
 )

 // PolynomialFit fits a polynomial of given degree to the data points (x, y).
-func PolynomialFit(x, y []float64, degree int) []float64 {
+func PolynomialFit(x, y []float64, degree int) ([]float64, error) {
 	n := len(x)
 	// Create the Vandermonde matrix
 	vander := mat.NewDense(n, degree+1, nil)
@@ -25,11 +23,7 @@ func PolynomialFit(x, y []float64, degree int) []float64 {
 	qr.Factorize(vander)
 	coeffs := mat.NewDense(degree+1, 1, nil)
 	err := qr.SolveTo(coeffs, false, yVec)
-	if err != nil {
-		panic(err)
-	}
-
-	return coeffs.RawMatrix().Data
+	return coeffs.RawMatrix().Data, err
 }

 // pow is a helper function to calculate power of a number.
@@ -43,19 +37,3 @@ func pow(x float64, n int) float64 {
 	}
 	return res
 }
-
-func main() {
-	// Input data
-	x := []float64{1.0, 2.0, 3.0, 4.0, 5.0}
-	y := []float64{1.0, 1.5, 2.5, 3.5, 5.0}
-
-	// Fit a second degree polynomial (quadratic)
-	degree := 2
-	coeffs := PolynomialFit(x, y, degree)
-
-	// Print the coefficients
-	fmt.Printf("Fitted polynomial coefficients: %v\n", coeffs)
-	if degree == 2 {
-		fmt.Printf("Fitted quadratic: y = %.4fx^2 + %.4fx + %.4f\n", coeffs[2], coeffs[1], coeffs[0])
-	}
-}
--- a/go.mod
+++ b/go.mod
@@ -9,14 +9,17 @@ require (
 )

 require (
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.17 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
 	golang.org/x/text v0.15.0 // indirect
 )

 require (
 	github.com/google/go-cmp v0.6.0 // indirect
 	github.com/k0kubun/pp/v3 v3.2.0
+	github.com/spf13/cobra v1.8.0
 	golang.org/x/net v0.25.0 // indirect
 	golang.org/x/sync v0.7.0 // indirect
 	golang.org/x/sys v0.20.0 // indirect
--- a/go.sum
+++ b/go.sum
@@ -835,6 +835,7 @@ github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4/go.mod h1:eXthEFrGJvWH
 github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/cpuguy83/go-md2man/v2 v2.0.0/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU=
+github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
@@ -1033,6 +1034,8 @@ github.com/iancoleman/strcase v0.2.0/go.mod h1:iwCmte+B7n89clKwxIoIXy/HfoL7AsD47
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
 github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
 github.com/json-iterator/go v1.1.11/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
@@ -1110,6 +1113,7 @@ github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFR
 github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
 github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/ruudk/golang-pdf417 v0.0.0-20181029194003-1af4ab5afa58/go.mod h1:6lfFZQK844Gfx8o5WFuvpxWRwnSoipWe/p622j1v06w=
 github.com/ruudk/golang-pdf417 v0.0.0-20201230142125-a7e3863a1245/go.mod h1:pQAZKsJ8yyVxGRWYNEm9oFB8ieLgKFnamEyDmSA0BRk=
 github.com/ryanuber/columnize v0.0.0-20160712163229-9b3edd62028f/go.mod h1:sm1tb6uqfes/u+d4ooFouqFdy9/2g9QGwK3SQygK0Ts=
@@ -1125,7 +1129,10 @@ github.com/spf13/afero v1.6.0/go.mod h1:Ai8FlHk4v/PARR026UzYexafAt9roJ7LcLMAmO6Z
 github.com/spf13/afero v1.9.2/go.mod h1:iUV7ddyEEZPO5gA3zD4fJt6iStLlL+Lg4m2cihcDf8Y=
 github.com/spf13/cast v1.3.1/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
 github.com/spf13/cobra v1.2.1/go.mod h1:ExllRjgxM/piMAM+3tAZvg8fsklGAf3tPfi+i8t68Nk=
+github.com/spf13/cobra v1.8.0 h1:7aJaZx1B85qltLMc546zn58BxxfZdR/W22ej9CFoEf0=
+github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
 github.com/spf13/jwalterweatherman v1.1.0/go.mod h1:aNWZUN0dPAAO/Ljvb5BEdw96iTZ0EXowPYD95IqWIGo=
+github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spf13/viper v1.8.1/go.mod h1:o0Pch8wJ9BVSWGQMbra6iw0oQ5oktSIBaujf1rJH9Ns=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
--- a/hdr.go
+++ b/hdr.go
@@ -0,0 +1,55 @@
+package imageproc
+
+import (
+	"bytes"
+	"html"
+	"os"
+
+	"text/template"
+)
+
+var HDRTemplate = `ENVI
+description = {
+  File Imported into ENVI.}
+samples = {{.Samples}}
+lines = {{.Lines}}
+bands = {{.Bands}}
+header offset = 0
+file type = ENVI Standard
+data type = 12
+interleave = bsq
+sensor type = Unknown
+byte order = 0
+wavelength units = Unknown`
+
+// https://www.nv5geospatialsoftware.com/docs/ENVIHeaderFiles.html
+type EnviHdr struct {
+	Samples int // samples (pixels) each line
+	Lines   int
+	Bands   int
+}
+
+func (h *EnviHdr) String() string {
+	var t *template.Template
+	t = template.New("template")
+	t, err := t.Parse(HDRTemplate)
+	if err != nil {
+		return ""
+	}
+	buf := new(bytes.Buffer)
+	t.Execute(buf, h)
+	buf = bytes.NewBufferString(html.UnescapeString(buf.String()))
+
+	return buf.String()
+}
+
+func (h *EnviHdr) Write(filename string) error {
+	str := h.String()
+	f, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	_, err = f.WriteString(str)
+	return err
+}
--- a/image_registration.go
+++ b/image_registration.go
@@ -1,7 +1,6 @@
 package imageproc

 import (
-	"bufio"
 	"fmt"
 	"image"
 	"image/color"
@@ -31,6 +30,8 @@ const (
 type ResampleMethod string

 type Registrator struct {
+	Params Params
+
 	PanImage  gocv.Mat
 	PanHeight int
 	PanWidth  int
@@ -50,15 +51,15 @@ type Registrator struct {
 	resampleMethod ResampleMethod
 }

-func NewRegistrator() *Registrator {
+func NewRegistrator(rsmethod ResampleMethod) *Registrator {
 	var r Registrator
-	r.resampleMethod = DownSampled
+	r.resampleMethod = rsmethod

 	return &r
 }

-func (r *Registrator) LoadPanRaw(raw string) error {
-	data, err := os.ReadFile(raw)
+func (r *Registrator) LoadPanRaw() error {
+	data, err := os.ReadFile(r.Params.PanRawFile)
 	if err != nil {
 		log.Error("Error reading raw file: ", err)
 		return err
@@ -92,8 +93,8 @@ func (r *Registrator) LoadPanRaw(raw string) error {
 	return nil
 }

-func (r *Registrator) LoadMssRaw(raw string) error {
-	data, err := os.ReadFile(raw)
+func (r *Registrator) LoadMssRaw() error {
+	data, err := os.ReadFile(r.Params.MssRawFile)
 	if err != nil {
 		log.Error("Error reading raw file: ", err)
 		return err
@@ -251,36 +252,7 @@ func (r *Registrator) calcPhaseCorrelation(panImage gocv.Mat,
 		}
 	}

-	r.calcDeltaCoeffs()
-
-	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.MssWidth * PixelBytes
-	height := r.MssHeight
-	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])
-	}
-
-	return nil
+	return r.calcDeltaCoeffs()
 }

 func (r *Registrator) Clean() {
@@ -321,8 +293,16 @@ func (r *Registrator) calcDeltaCoeffs() error {

 		}

-		r.deltaXCoeffs[i] = PolynomialFit(cx, dx, 1)
-		r.deltaYCoeffs[i] = PolynomialFit(cx, dy, 2)
+		var err error
+		if r.deltaXCoeffs[i], err = PolynomialFit(cx, dx, 1); err != nil {
+			log.Error("Error fitting deltaX coeffs: ", err)
+			return err
+		}
+
+		if r.deltaYCoeffs[i], err = PolynomialFit(cx, dy, 2); err != nil {
+			log.Error("Error fitting deltaY coeffs: ", err)
+			return err
+		}
 	}

 	for i := 0; i < MssBands; i++ {
--- a/output.go
+++ b/output.go
@@ -2,7 +2,10 @@ package imageproc

 import (
 	"bufio"
+	"fmt"
 	"os"
+	"path/filepath"
+	"strings"

 	"github.com/airbusgeo/godal"
 	log "github.com/sirupsen/logrus"
@@ -151,3 +154,41 @@ func (r *Registrator) BytesToRaw(mssData []byte, filePath string) error {

 	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.MssWidth * PixelBytes
+	height := r.MssHeight
+	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: r.MssWidth,
+		Lines:   r.MssHeight,
+		Bands:   MssBands,
+	}
+
+	hdrFile := filepath.Join(filepath.Dir(raw),
+		fmt.Sprintf("%s.HDR", strings.TrimSuffix(filepath.Base(raw), filepath.Ext(raw))),
+	)
+	hdr.Write(hdrFile)
+
+	return nil
+}
--- a/params.go
+++ b/params.go
@@ -0,0 +1,13 @@
+package imageproc
+
+type Params struct {
+	PanRawFile           string
+	MssRawFile           string
+	AuxRawFile           string
+	SaveRegisteredMssRaw bool
+	PansharpenIHS        bool
+	OutputDir            string
+	PanTiffFile          string
+	MssTiffFile          string
+	FusTIffFile          string
+}