EasyVQD/pkg/vqdcms/vqd.go

package vqdcms

/*
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
*/
import "C"
import (
	"easyvqd/pkg/decoder"
	"fmt"
	"github.com/shirou/gopsutil/v4/mem"
	"log/slog"
	"os"
	"path/filepath"
	"runtime/debug"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unsafe"
)

const MAX_STREAM_CHAN_NUM = 256

type VideoInfoVQD struct {
	mu              sync.RWMutex
	VQDHandle       uintptr
	Params          VQDPara
	IsCreateSuccess bool
}

type VQDHandle struct {
	running  uint32
	TaskID   int    // 标识ID
	ChnID    string // 通道ID
	Plans    string // 任务计划
	data     chan ChanData
	dataLock sync.RWMutex
	cb       VQDResultCB
	handle   *VideoInfoVQD
	name     string // 算法名称
	file     *os.File
	decoder  *decoder.VideoDecoder
	fileLock sync.RWMutex
}

type VQDResultCB func(AbnormalModel)

type ChanData struct {
	data []byte
	w    int
	h    int
	now  time.Time
}

func IsCurTimeInRecordPlan(recordPlanNew string, now time.Time) bool {
	return false
}

func NewVQDHandle(cb VQDResultCB, taskId int, chnId string) *VQDHandle {
	v := &VQDHandle{
		running: 0,
		decoder: &decoder.VideoDecoder{},
		ChnID:   chnId,
		TaskID:  taskId,
		data:    make(chan ChanData, MAX_STREAM_CHAN_NUM),
		cb:      cb,
		handle:  &VideoInfoVQD{},
	}
	err := v.decoder.Create()
	if err != nil {
		slog.Error("decoder Create ", "taskId", chnId, "err", err)
	}
	return v
}

func (v *VQDHandle) SetVQDConfig(params VQDPara) error {
	return v.handle.Config(params, params.EnableFunc, 10)
}

func (v *VQDHandle) GetHandle() *VideoInfoVQD {
	return v.handle
}
func (v *VQDHandle) Create(params VQDPara) *VQDHandle {
	if atomic.LoadUint32(&v.running) == 1 {
		return v
	}
	//if v.file == nil {
	//	filename := filepath.Join(utils.CWD(), "logs", fmt.Sprintf("%016p_vqdhandle.log", v))
	//	file, err := os.OpenFile(filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, os.ModePerm)
	//	if err == nil {
	//		v.fileLock.Lock()
	//		v.file = file
	//		v.file.WriteString("↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓\r\n")
	//		v.fileLock.Unlock()
	//	} else {
	//		llog.Info("创建文件失败：%s", filename)
	//	}
	//}
	err := v.handle.Create(params, params.EnableFunc, 10)
	if err != nil {
		//llog.Info("vqd create fail:%s", err.Error())
		return v
	}
	atomic.StoreUint32(&v.running, 1)
	go v.RunFrame()
	return v
}

func (v *VQDHandle) Destroy() {
	if v.file != nil {
		v.fileLock.Lock()
		v.file.WriteString("↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑\r\n")
		v.file.Close()
		v.file = nil
		v.fileLock.Unlock()
	}
	if v.data != nil {
		close(v.data)
	}
	v.data = nil
	if atomic.LoadUint32(&v.running) == 1 {
		atomic.StoreUint32(&v.running, 0)
		v.handle.Destroy()
	}
}

func (v *VQDHandle) RunFrame() {
	defer func() {
		if e := recover(); e != nil {
			print(fmt.Sprintf("RunFrame---%s\n", e))
			print(fmt.Sprintf("%s\n", string(debug.Stack())))
		}
	}()
	cvqdImgsDir := filepath.Join(CWD(), "vqd_images", fmt.Sprintf("%d", v.TaskID))
	cvqdImgsDir = filepath.ToSlash(cvqdImgsDir)
	if err := os.MkdirAll(cvqdImgsDir, os.ModePerm); err != nil {
		//llog.Info("%s:%s", cvqdImgsDir, err.Error())
	}

	hyper := 0
	index := 0
	cdata := ChanData{}
	for {
		select {
		case data, ok := <-v.data:
			if !ok {
				return
			}
			if len(v.data) >= (MAX_STREAM_CHAN_NUM - 6) {
				//llog.Info("vqd channel num >= %d", MAX_STREAM_CHAN_NUM)
				hyper = MAX_STREAM_CHAN_NUM / 10
			}
			if hyper > 0 {
				hyper--
				break
			}
			cdata = data
			//if !IsCurTimeInRecordPlan(v.analysisTime, cdata.now) {
			//	continue
			//}
			if v.file != nil {
				v.fileLock.Lock()
				v.file.WriteString(fmt.Sprintf("vqd START data[%d] w[%d] h[%d]\r\n", len(cdata.data), cdata.w, cdata.h))
				v.fileLock.Unlock()
			}
			now := time.Now().UnixMilli()
			fpath := filepath.Join(cvqdImgsDir, fmt.Sprintf("%d_%s_%d.jpg", v.TaskID, v.ChnID, now))
			fpath = filepath.ToSlash(fpath)
			result := VQDResult{}
			ret := v.handle.Frame(cdata.data, cdata.w, cdata.h, index, fpath, &result)
			if ret == 0 {
				index = 0
				if value, b := v.parseVQD(result); b {
					value.FilePath = strings.TrimPrefix(filepath.ToSlash(fpath), filepath.ToSlash(CWD()))
					value.ChannelID = v.ChnID
					value.AlgoName = v.name
					if v.cb != nil {
						v.cb(value)
					} else {
						// 保存数据库
						//CreateAbnormalModel(&value)
					}
				}
			}
			if v.file != nil {
				me, err := mem.VirtualMemory()
				if err != nil {
					v.fileLock.Lock()
					v.file.WriteString(fmt.Sprintf("vqd END ret[%d] mem is nil\r\n", ret))
					v.fileLock.Unlock()
				} else {
					v.fileLock.Lock()
					v.file.WriteString(fmt.Sprintf("vqd END ret[%d] mem[%.2f%s]\r\n", ret, me.UsedPercent, "%"))
					v.fileLock.Unlock()
				}
			}
			//C.free(fp)
			//C.free(yuvBuf)
			index++
		}
	}
}

func (v *VQDHandle) parseVQD(result VQDResult) (AbnormalModel, bool) {
	isabnormal := false
	abnormals := AbnormalModel{
		IsDeep:    v.handle.Params.UseDeepLearning,
		Abnormals: make([]Abnormal, 0),
	}

	if (result.AbnormalType & NXU_VQD_ABN_COLORDEV) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.ColorDev,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_COLORDEV],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.ColorPara.ColorThr,
		Name1: "偏色阈值",
		Ratio: v.handle.Params.ColorPara.ColorAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_LIGHT) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.LgtDark,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_LIGHT],
		})
	}

	if (result.AbnormalType & NXU_VQD_ABN_DARK) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.LgtDark,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_DARK],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.LgtDarkPara.LightThr,
		Name1: "过亮阈值",
		Thr2:  v.handle.Params.LgtDarkPara.DarkThr,
		Name2: "过暗阈值",
		Ratio: v.handle.Params.LgtDarkPara.LgtDarkAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_CLARITY) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Clarity,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_CLARITY],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.ClarityPara.ClarityThr,
		Name1: "清晰度阈值",
		Ratio: v.handle.Params.ClarityPara.ClarityAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_NOISE) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Noise,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_NOISE],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.NoisePara.NoiseThr,
		Name1: "噪声阈值",
		Ratio: v.handle.Params.NoisePara.NoiseAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_CONTRAST) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Contrast,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_CONTRAST],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.ContrastPara.CtraLowThr,
		Name1: "低对比度阈值",
		Thr2:  v.handle.Params.ContrastPara.CtraHighThr,
		Name2: "高对比度阈值",
		Ratio: v.handle.Params.ContrastPara.CtraAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_OCCLUSION) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Occlusion,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_OCCLUSION],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.OcclusionPara.OcclusionThr,
		Name1: "遮挡阈值",
		Ratio: v.handle.Params.OcclusionPara.OcclusionAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_BLUE) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Blue,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_BLUE],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.BluePara.BlueThr,
		Name1: "蓝屏阈值",
		Ratio: v.handle.Params.BluePara.BlueAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_SHARK) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Shark,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_SHARK],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.SharkPara.SharkThr,
		Name1: "抖动阈值",
		Ratio: v.handle.Params.SharkPara.SharkAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_FREEZE) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Freeze,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_FREEZE],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.FreezePara.FreezeThr,
		Name1: "冻结阈值",
		Ratio: v.handle.Params.FreezePara.FreezeAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_MOSAIC) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Mosaic,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_MOSAIC],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.MosaicPara.MosaicThr,
		Name1: "马赛克阈值",
		Ratio: v.handle.Params.MosaicPara.MosaicAbnNumRatio,
	})

	if (result.AbnormalType & NXU_VQD_ABN_FLOWER) != 0 {
		isabnormal = true
		abnormals.Abnormals = append(abnormals.Abnormals, Abnormal{
			Value: result.Flower,
			Name:  ALNORMAL_NAMES[NXU_VQD_ABN_FLOWER],
		})
	}
	abnormals.DefaultValues = append(abnormals.DefaultValues, DefaultValue{
		Thr1:  v.handle.Params.FlowerPara.FlowerThr,
		Name1: "花屏阈值",
		Ratio: v.handle.Params.FlowerPara.FlowerAbnNumRatio,
	})

	return abnormals, isabnormal
}
func (v *VQDHandle) SendData(buf []byte, _codec int) {
	dataP := GetIFramePointer(buf)
	if dataP == nil || !dataP.IsValid {
		slog.Error("I帧转指针失败: ", "TaskID", v.TaskID)
		return
	}
	w, h, data, err := v.decoder.PushDataEx(uintptr(dataP.Pointer), dataP.Length, _codec)
	if err != nil {
		slog.Error("I帧转YUV失败: ", "TaskID", v.TaskID, "err", err)
		return
	}
	if len(data) > 0 && v.data != nil {
		dst := make([]byte, len(data))
		copy(dst, data)
		now := time.Now()
		d := ChanData{
			data: dst,
			w:    w,
			h:    h,
			now:  now,
		}
		//v.dataLock.Lock()
		//v.data = append(v.data, d)
		//v.dataLock.Unlock()
		v.data <- d
	}
}

// IFrameData 封装I帧数据和指针信息
type IFrameData struct {
	Data    []byte         // I帧原始字节数据
	Pointer unsafe.Pointer // 指向数据的原始指针
	Length  int            // 数据长度（字节数）
	IsValid bool           // 指针是否有效
}

// GetIFramePointer 将字节切片转换为原始指针
// 注意：unsafe包的使用会绕过Go的内存安全检查，需谨慎
func GetIFramePointer(data []byte) *IFrameData {
	if len(data) == 0 {
		return &IFrameData{
			IsValid: false,
			Length:  0,
		}
	}

	// 方式1：直接通过unsafe获取切片底层数组的指针（推荐，高效）
	// 切片的底层结构是：指向数组的指针 + 长度 + 容量
	ptr := unsafe.Pointer(&data[0])

	// 方式2：通过reflect获取指针（更直观展示切片结构，可选）
	// sliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	// ptr := unsafe.Pointer(sliceHeader.Data)

	return &IFrameData{
		Data:    data,
		Pointer: ptr,
		Length:  len(data),
		IsValid: true,
	}
}
func CWD() string {
	geTwd, err := os.Getwd()
	if err != nil {
		return ""
	}
	return geTwd
}