feat: option to toggle contrast enhancement in 2dfan4

model/2dfan4/2dfan4.go

@@ -12,6 +12,17 @@ type Model struct {
 type Input struct {
 	Img         gocv.Mat
 	BoundingBox model.BoundingBox
+	// ApplyContrastEnhancement allows toggling this feature on/off.
+	// NOTE: Enabling this will slow down processing because it performs extra color conversions and CLAHE.
+	ApplyContrastEnhancement bool
+	// BrightnessThreshold defines how dark the image can be before we apply contrast enhancement.
+	// Suggested value: 30.0 for low-light detection.
+	// Typical LAB L-channel values range from 0 to 100.
+	BrightnessThreshold float64
+	// ClaheClipLimit controls how much CLAHE enhances contrast.
+	// Suggested value: 2.0 for a balanced effect without over-enhancement.
+	// Higher values = more aggressive contrast stretching. Typical values: 2.0 to 4.0.
+	ClaheClipLimit float32
 }
 
 type Output struct {

model/2dfan4/pre.go

@@ -17,6 +17,47 @@ func (m *Model) PreProcess(i *Input) ([]*protobuf.InferTensorContents, error) {
 	cropVisionFrame, affineMatrix := model.WarpFaceByTranslation(i.Img, translation, scale, image.Point{X: 256, Y: 256})
 	defer cropVisionFrame.Close()
 	m.affineMatrix = affineMatrix
+
+	if i.ApplyContrastEnhancement {
+		applyContrastEnhancement(cropVisionFrame, i.BrightnessThreshold, i.ClaheClipLimit)
+	}
+
+	cropVisionFrame.ConvertTo(&cropVisionFrame, gocv.MatTypeCV32F)
+	cropVisionFrame.DivideFloat(255.0)
+
+	rgbChannels := gocv.Split(cropVisionFrame)
+	defer rgbChannels[0].Close()
+	defer rgbChannels[1].Close()
+	defer rgbChannels[2].Close()
+
+	totalSize := rgbChannels[0].Rows() * rgbChannels[0].Cols()
+	d := make([]float32, 3*totalSize)
+
+	rd, err := rgbChannels[0].DataPtrFloat32()
+	if err != nil {
+		return nil, err
+	}
+	copy(d[0:totalSize], rd)
+
+	gd, err := rgbChannels[1].DataPtrFloat32()
+	if err != nil {
+		return nil, err
+	}
+	copy(d[totalSize:2*totalSize], gd)
+
+	bd, err := rgbChannels[2].DataPtrFloat32()
+	if err != nil {
+		return nil, err
+	}
+	copy(d[2*totalSize:3*totalSize], bd)
+
+	contents := &protobuf.InferTensorContents{
+		Fp32Contents: d,
+	}
+	return []*protobuf.InferTensorContents{contents}, nil
+}
+
+func applyContrastEnhancement(cropVisionFrame gocv.Mat, threshold float64, clipLimit float32) {
 	gocv.CvtColor(cropVisionFrame, &cropVisionFrame, gocv.ColorRGBToLab)
 
 	labChannels := gocv.Split(cropVisionFrame)
@@ -25,39 +66,12 @@ func (m *Model) PreProcess(i *Input) ([]*protobuf.InferTensorContents, error) {
 	defer labChannels[2].Close()
 
 	meanValue := calculateMean(labChannels[0])
-	threshold := 30.0
-	clipLimit := float32(2.0)
 
 	if meanValue < threshold {
 		applyCLAHE(labChannels[0], clipLimit)
 	}
 
-	d := []float32{}
-
 	gocv.CvtColor(cropVisionFrame, &cropVisionFrame, gocv.ColorLabToRGB)
-
-	cropVisionFrame.ConvertTo(&cropVisionFrame, gocv.MatTypeCV32F)
-	cropVisionFrame.DivideFloat(255.0)
-	rgbChannels := gocv.Split(cropVisionFrame)
-	r := rgbChannels[0]
-	defer r.Close()
-	rd, _ := r.DataPtrFloat32()
-	d = append(d, rd...)
-
-	g := rgbChannels[1]
-	defer g.Close()
-	gd, _ := g.DataPtrFloat32()
-	d = append(d, gd...)
-
-	b := rgbChannels[2]
-	defer b.Close()
-	bd, _ := b.DataPtrFloat32()
-	d = append(d, bd...)
-
-	contents := &protobuf.InferTensorContents{
-		Fp32Contents: d,
-	}
-	return []*protobuf.InferTensorContents{contents}, nil
 }
 
 // calculateScale calculates the scale factor based on the bounding box dimensions.