convolution.go 2.9 KB

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