pre_process.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566
  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. def get_classes(classes_path):
  28. """loads the classes"""
  29. with open(classes_path) as f:
  30. class_names = f.readlines()
  31. class_names = [c.strip() for c in class_names]
  32. return class_names
  33. def get_anchors(anchors_path):
  34. """loads the anchors from a file"""
  35. with open(anchors_path) as f:
  36. anchors = f.readline()
  37. anchors = [float(x) for x in anchors.split(',')]
  38. return np.array(anchors).reshape(-1, 2)
  39. def get_colors(number, bright=True):
  40. """
  41. Generate random colors for drawing bounding boxes.
  42. To get visually distinct colors, generate them in HSV space then
  43. convert to RGB.
  44. """
  45. if number <= 0:
  46. return []
  47. brightness = 1.0 if bright else 0.7
  48. hsv_tuples = [(x / number, 1., brightness)
  49. for x in range(number)]
  50. colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
  51. colors = list(
  52. map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
  53. colors))
  54. np.random.seed(10101) # Fixed seed for consistent colors across runs.
  55. np.random.shuffle(colors) # Shuffle colors to decorrelate adjacent classes.
  56. np.random.seed(None) # Reset seed to default.
  57. return colors