git.sr.ht/~pingoo/stdx@v0.0.0-20240218134121-094174641f6e/imaging/convolution.go (about)

     1  package imaging
     2  
     3  import (
     4  	"image"
     5  )
     6  
     7  // ConvolveOptions are convolution parameters.
     8  type ConvolveOptions struct {
     9  	// If Normalize is true the kernel is normalized before convolution.
    10  	Normalize bool
    11  
    12  	// If Abs is true the absolute value of each color channel is taken after convolution.
    13  	Abs bool
    14  
    15  	// Bias is added to each color channel value after convolution.
    16  	Bias int
    17  }
    18  
    19  // Convolve3x3 convolves the image with the specified 3x3 convolution kernel.
    20  // Default parameters are used if a nil *ConvolveOptions is passed.
    21  func Convolve3x3(img image.Image, kernel [9]float64, options *ConvolveOptions) *image.NRGBA {
    22  	return convolve(img, kernel[:], options)
    23  }
    24  
    25  // Convolve5x5 convolves the image with the specified 5x5 convolution kernel.
    26  // Default parameters are used if a nil *ConvolveOptions is passed.
    27  func Convolve5x5(img image.Image, kernel [25]float64, options *ConvolveOptions) *image.NRGBA {
    28  	return convolve(img, kernel[:], options)
    29  }
    30  
    31  func convolve(img image.Image, kernel []float64, options *ConvolveOptions) *image.NRGBA {
    32  	src := toNRGBA(img)
    33  	w := src.Bounds().Max.X
    34  	h := src.Bounds().Max.Y
    35  	dst := image.NewNRGBA(image.Rect(0, 0, w, h))
    36  
    37  	if w < 1 || h < 1 {
    38  		return dst
    39  	}
    40  
    41  	if options == nil {
    42  		options = &ConvolveOptions{}
    43  	}
    44  
    45  	if options.Normalize {
    46  		normalizeKernel(kernel)
    47  	}
    48  
    49  	type coef struct {
    50  		x, y int
    51  		k    float64
    52  	}
    53  	var coefs []coef
    54  	var m int
    55  
    56  	switch len(kernel) {
    57  	case 9:
    58  		m = 1
    59  	case 25:
    60  		m = 2
    61  	}
    62  
    63  	i := 0
    64  	for y := -m; y <= m; y++ {
    65  		for x := -m; x <= m; x++ {
    66  			if kernel[i] != 0 {
    67  				coefs = append(coefs, coef{x: x, y: y, k: kernel[i]})
    68  			}
    69  			i++
    70  		}
    71  	}
    72  
    73  	parallel(0, h, func(ys <-chan int) {
    74  		for y := range ys {
    75  			for x := 0; x < w; x++ {
    76  				var r, g, b float64
    77  				for _, c := range coefs {
    78  					ix := x + c.x
    79  					if ix < 0 {
    80  						ix = 0
    81  					} else if ix >= w {
    82  						ix = w - 1
    83  					}
    84  
    85  					iy := y + c.y
    86  					if iy < 0 {
    87  						iy = 0
    88  					} else if iy >= h {
    89  						iy = h - 1
    90  					}
    91  
    92  					off := iy*src.Stride + ix*4
    93  					s := src.Pix[off : off+3 : off+3]
    94  					r += float64(s[0]) * c.k
    95  					g += float64(s[1]) * c.k
    96  					b += float64(s[2]) * c.k
    97  				}
    98  
    99  				if options.Abs {
   100  					if r < 0 {
   101  						r = -r
   102  					}
   103  					if g < 0 {
   104  						g = -g
   105  					}
   106  					if b < 0 {
   107  						b = -b
   108  					}
   109  				}
   110  
   111  				if options.Bias != 0 {
   112  					r += float64(options.Bias)
   113  					g += float64(options.Bias)
   114  					b += float64(options.Bias)
   115  				}
   116  
   117  				srcOff := y*src.Stride + x*4
   118  				dstOff := y*dst.Stride + x*4
   119  				d := dst.Pix[dstOff : dstOff+4 : dstOff+4]
   120  				d[0] = clamp(r)
   121  				d[1] = clamp(g)
   122  				d[2] = clamp(b)
   123  				d[3] = src.Pix[srcOff+3]
   124  			}
   125  		}
   126  	})
   127  
   128  	return dst
   129  }
   130  
   131  func normalizeKernel(kernel []float64) {
   132  	var sum, sumpos float64
   133  	for i := range kernel {
   134  		sum += kernel[i]
   135  		if kernel[i] > 0 {
   136  			sumpos += kernel[i]
   137  		}
   138  	}
   139  	if sum != 0 {
   140  		for i := range kernel {
   141  			kernel[i] /= sum
   142  		}
   143  	} else if sumpos != 0 {
   144  		for i := range kernel {
   145  			kernel[i] /= sumpos
   146  		}
   147  	}
   148  }