Guest User

color recognition picamera

a guest
Apr 6th, 2017
118
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import picamera
  2. import picamera.array
  3. import time
  4. import cv2
  5.  
  6. camSize = 320, 240
  7. sampling = 4 # lower value is more sampling ie 1 samples all
  8. minCandidates = 20 # minimun number of valid pixels for a track
  9. searchDimension = 30
  10. crosshairSize = 30
  11. colorStrengthThreshold = 2
  12. colorArray = [ [[255, 0, 0], [255, 0, 0], "red"], [[0, 255, 0], [0, 255, 0], "green"], [[0, 0, 255], [0, 0, 255], "blue"], [[0, 255, 255], [0, 255, 255], "cyan"], [[255, 255, 0], [255, 255, 0], "yellow"], [[255, 0, 255], [255, 0, 255], "magenta"], [[0, 0, 0], [0, 0, 0], "background"] ] # r, g, b, c, y, m, default
  13. colorNames = ["red", "green", "blue", "cyan", "yellow", "magenta", "default"]
  14. colorDistanceThreshold = 30000
  15.  
  16. trainPrompt = raw_input("Trainer mode? (y/n): ")
  17. if trainPrompt.lower() == "y":
  18. for i in range(0, len(colorArray)):
  19. z = raw_input("train " + colorNames[i] + ": ")
  20. if z == "y":
  21. trainAvgColor = [0, 0, 0]
  22. # TODO: All of this can just one one initialized camera - we don't need to do this twice
  23. with picamera.PiCamera() as trainCam:
  24. trainCam.awb_mode = 'off'
  25. trainCam.awb_gains = 1.5
  26. trainCam.exposure_mode = 'fixedfps'
  27. trainCam.resolution = (camSize[0], camSize[1])
  28. time.sleep(0.5)
  29. with picamera.array.PiRGBArray(trainCam) as trainStream:
  30. trainCam.capture(trainStream, resize = (camSize[0], camSize[1]), format = 'rgb')
  31. trainImage = trainStream.array
  32. trainPixelReadCount = 0
  33. for y in range(camSize[1] / 2 - searchDimension, camSize[1] / 2 + searchDimension):
  34. for x in range(camSize[0] / 2 - searchDimension, camSize[0] / 2 + searchDimension):
  35. if x % sampling == 0 and y % sampling == 0:
  36. trainColor = trainImage[y][x]
  37.  
  38. trainPixelReadCount += 1
  39. trainAvgColor[0] += trainColor[0]
  40. trainAvgColor[1] += trainColor[1]
  41. trainAvgColor[2] += trainColor[2]
  42.  
  43. trainAvgColor = trainAvgColor[0] / trainPixelReadCount, trainAvgColor[1] / trainPixelReadCount, trainAvgColor[2] / trainPixelReadCount
  44.  
  45. colorArray[i][1] = trainAvgColor
  46.  
  47.  
  48. while True:
  49. with picamera.PiCamera() as cam:
  50. cam.awb_mode = 'off'
  51. cam.awb_gains = 1.5
  52. cam.exposure_mode = 'fixedfps'
  53.  
  54. cam.resolution = (camSize[0], camSize[1])
  55. cam.start_preview()
  56. time.sleep(1)
  57. with picamera.array.PiRGBArray(cam) as stream:
  58. cam.capture(stream, resize = (camSize[0], camSize[1]), format = 'rgb')
  59. image = stream.array
  60.  
  61. avgColor = [0, 0, 0]
  62. pixelReadCount = 0
  63.  
  64. for y in range(camSize[1] / 2 - searchDimension, camSize[1] / 2 + searchDimension):
  65. for x in range(camSize[0] / 2 - searchDimension, camSize[0] / 2 + searchDimension):
  66. if x % sampling == 0 and y % sampling == 0:
  67. color = image[y][x]
  68.  
  69. pixelReadCount += 1
  70. avgColor[0] += color[0]
  71. avgColor[1] += color[1]
  72. avgColor[2] += color[2]
  73.  
  74. avgColor = avgColor[0] / pixelReadCount, avgColor[1] / pixelReadCount, avgColor[2] / pixelReadCount
  75.  
  76. lastDist = colorDistanceThreshold
  77. selectedColor = colorArray[6]
  78.  
  79. for c in colorArray:
  80. colorDist = pow(avgColor[0] - c[0][0], 2) + pow(avgColor[1] - c[0][1], 2) + pow(avgColor[2] - c[0][2], 2)
  81. if colorDist < lastDist:
  82. lastDist = colorDist
  83. selectedColor = c
  84.  
  85. print "distance: " + str(lastDist)
  86. print "selected color: " + str(selectedColor[2])
  87.  
  88.  
  89. time.sleep(0.5)
RAW Paste Data