pre_process.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. import colorsys
  2. import time
  3. import numpy as np
  4. import cv2
  5. from skimage import measure
  6. def count_red_pixel(image_np, cnt=1000):
  7. # 红色像素计数
  8. start_time = time.time()
  9. image_hsv = cv2.cvtColor(image_np, cv2.COLOR_BGR2HSV)
  10. # minus_1 = image_np[:, :, 2].astype('int32') - image_np[:, :, 0].astype('int32')
  11. # minus_2 = image_np[:, :, 2].astype('int32') - image_np[:, :, 1].astype('int32')
  12. # red_mask = (image_np[:, :, 2] >= 180) & (minus_1 >= 60) & (minus_2 >= 60)
  13. red_mask = ((image_hsv[:, :, 0] >= 0) & (image_hsv[:, :, 0] <= 10) | (image_hsv[:, :, 0] <= 180) & (image_hsv[:, :, 0] >= 156)) \
  14. & (image_hsv[:, :, 1] <= 255) & (image_hsv[:, :, 1] >= 43) \
  15. & (image_hsv[:, :, 2] <= 255) & (image_hsv[:, :, 2] >= 100)
  16. red_cnt = 0
  17. labels = measure.label(red_mask, connectivity=2) # 8连通区域标记
  18. regions = measure.regionprops(labels)
  19. red_cnt = np.sum(red_mask != 0)
  20. print("red_cnt regions", len(regions),red_cnt, time.time()-start_time)
  21. if regions and len(regions)>0:
  22. _max_area = max([r.bbox_area for r in regions])
  23. if _max_area>100:
  24. print("red_cnt max_area", _max_area, time.time()-start_time)
  25. return True
  26. return False
  27. if red_cnt >= cnt:
  28. return True
  29. else:
  30. return False
  31. def get_classes(classes_path):
  32. """loads the classes"""
  33. with open(classes_path) as f:
  34. class_names = f.readlines()
  35. class_names = [c.strip() for c in class_names]
  36. return class_names
  37. def get_anchors(anchors_path):
  38. """loads the anchors from a file"""
  39. with open(anchors_path) as f:
  40. anchors = f.readline()
  41. anchors = [float(x) for x in anchors.split(',')]
  42. return np.array(anchors).reshape(-1, 2)
  43. def get_colors(number, bright=True):
  44. """
  45. Generate random colors for drawing bounding boxes.
  46. To get visually distinct colors, generate them in HSV space then
  47. convert to RGB.
  48. """
  49. if number <= 0:
  50. return []
  51. brightness = 1.0 if bright else 0.7
  52. hsv_tuples = [(x / number, 1., brightness)
  53. for x in range(number)]
  54. colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
  55. colors = list(
  56. map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
  57. colors))
  58. np.random.seed(10101) # Fixed seed for consistent colors across runs.
  59. np.random.shuffle(colors) # Shuffle colors to decorrelate adjacent classes.
  60. np.random.seed(None) # Reset seed to default.
  61. return colors