package main

import (
	"fmt"
	"gocv.io/x/gocv"
	"image"
	"io/ioutil"
	"log"
	"time"
)

func main(){

	start := time.Now()

	//加载模型
	model := gocv.ReadNetFromONNX("/Users/ljq/code/weight/ShipWeight_6-25_0.96.onnx")
	modelLoadtime := time.Since(start)
	fmt.Printf("modelLoadtime====[%s]\n",modelLoadtime)

	Imgtime := time.Now()

	dirQ := "val/Q"
	fileinfo, err := ioutil.ReadDir(dirQ)
	if err != nil {
		fmt.Printf("e",err)
	}

	T := 0
	F := 0

	//遍历Q测试集
	for _, fi := range fileinfo {
		filename := fi.Name()
		filename = dirQ + "/"+ filename
		//fmt.Printf("Q",filename,"\n")
		Shipimg := gocv.IMRead(filename, gocv.IMReadColor)
		if Shipimg.Empty() {
			log.Fatal("Invalid Mat")
		}
		//图像预处理
		Shipimg = gocv.BlobFromImage(Shipimg, 1 / 255.0, image.Point{224, 224}, gocv.Scalar{},true,false)
		//
		model.SetInput(Shipimg,"input")
		//输出的是Mat对象
		out := model.Forward("output")

		var shipout []float32
		shipout,_ = out.DataPtrFloat32()
		//fmt.Printf("out",shipout,"\n")

		if shipout[0] > shipout[1]{
			fmt.Printf("这是轻载\n")
			T += 1

		}else {
			fmt.Printf("这是重载船\n")
			F += 1
		}
	}

	dirZ := "val/Z"
	fileinfoZ, err := ioutil.ReadDir(dirZ)
	if err != nil {
		fmt.Printf("e",err)
	}

	for _, fi := range fileinfoZ {
		filename := fi.Name()
		filename = dirZ + "/"+ filename
		//fmt.Printf("Q",filename,"\n")
		Shipimg := gocv.IMRead(filename, gocv.IMReadColor)
		if Shipimg.Empty() {
			log.Fatal("Invalid Mat")
		}
		//图像预处理
		Shipimg = gocv.BlobFromImage(Shipimg, 1 / 255.0, image.Point{224, 224}, gocv.Scalar{},true,false)
		//预设输入名
		model.SetInput(Shipimg,"input")
		//输出的是Mat对象
		out := model.Forward("output")

		var shipout []float32
		shipout,_ = out.DataPtrFloat32()
		//fmt.Printf("out",shipout,"\n")

		if shipout[0] > shipout[1]{
			fmt.Printf("这是轻载\n")
			F += 1

		}else {
			fmt.Printf("这是重载船\n")
			T += 1
		}
	}

	fmt.Printf("准确率：-----%f-----\n",T/(T+F))

	cost := time.Since(Imgtime)
	fmt.Printf("cost=[%s]---Per-----[%s]",cost,(T+F))

}