kmeans.py 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  1. #!/usr/bin/env python
  2. '''
  3. K-means clusterization sample.
  4. Usage:
  5. kmeans.py
  6. Keyboard shortcuts:
  7. ESC - exit
  8. space - generate new distribution
  9. '''
  10. # Python 2/3 compatibility
  11. from __future__ import print_function
  12. import numpy as np
  13. import cv2
  14. from gaussian_mix import make_gaussians
  15. if __name__ == '__main__':
  16. cluster_n = 5
  17. img_size = 512
  18. print(__doc__)
  19. # generating bright palette
  20. colors = np.zeros((1, cluster_n, 3), np.uint8)
  21. colors[0,:] = 255
  22. colors[0,:,0] = np.arange(0, 180, 180.0/cluster_n)
  23. colors = cv2.cvtColor(colors, cv2.COLOR_HSV2BGR)[0]
  24. while True:
  25. print('sampling distributions...')
  26. points, _ = make_gaussians(cluster_n, img_size)
  27. term_crit = (cv2.TERM_CRITERIA_EPS, 30, 0.1)
  28. ret, labels, centers = cv2.kmeans(points, cluster_n, None, term_crit, 10, 0)
  29. img = np.zeros((img_size, img_size, 3), np.uint8)
  30. for (x, y), label in zip(np.int32(points), labels.ravel()):
  31. c = list(map(int, colors[label]))
  32. cv2.circle(img, (x, y), 1, c, -1)
  33. cv2.imshow('gaussian mixture', img)
  34. ch = 0xFF & cv2.waitKey(0)
  35. if ch == 27:
  36. break
  37. cv2.destroyAllWindows()