高速公路违停检测
Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

119 Zeilen
3.6KB

  1. #!/usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. from PIL import Image
  4. import PIL.ImageEnhance as ImageEnhance
  5. import random
  6. import numpy as np
  7. class RandomCrop(object):
  8. def __init__(self, size, *args, **kwargs):
  9. self.size = size
  10. def __call__(self, im_lb):
  11. im = im_lb['im']
  12. lb = im_lb['lb']
  13. assert im.size == lb.size
  14. W, H = self.size
  15. w, h = im.size
  16. if (W, H) == (w, h): return dict(im=im, lb=lb)
  17. if w < W or h < H:
  18. scale = float(W) / w if w < h else float(H) / h
  19. w, h = int(scale * w + 1), int(scale * h + 1)
  20. im = im.resize((w, h), Image.BILINEAR)
  21. lb = lb.resize((w, h), Image.NEAREST)
  22. sw, sh = random.random() * (w - W), random.random() * (h - H)
  23. crop = int(sw), int(sh), int(sw) + W, int(sh) + H
  24. return dict(
  25. im = im.crop(crop),
  26. lb = lb.crop(crop)
  27. )
  28. class HorizontalFlip(object):
  29. def __init__(self, p=0.5, *args, **kwargs):
  30. self.p = p
  31. def __call__(self, im_lb):
  32. if random.random() > self.p:
  33. return im_lb
  34. else:
  35. im = im_lb['im']
  36. lb = im_lb['lb']
  37. return dict(im = im.transpose(Image.FLIP_LEFT_RIGHT),
  38. lb = lb.transpose(Image.FLIP_LEFT_RIGHT),
  39. )
  40. class RandomScale(object):
  41. def __init__(self, scales=(1, ), *args, **kwargs):
  42. self.scales = scales
  43. # print('scales: ', scales)
  44. def __call__(self, im_lb):
  45. im = im_lb['im']
  46. lb = im_lb['lb']
  47. W, H = im.size
  48. scale = random.choice(self.scales)
  49. # scale = np.random.uniform(min(self.scales), max(self.scales))
  50. w, h = int(W * scale), int(H * scale)
  51. return dict(im = im.resize((w, h), Image.BILINEAR),
  52. lb = lb.resize((w, h), Image.NEAREST),
  53. )
  54. class ColorJitter(object):
  55. def __init__(self, brightness=None, contrast=None, saturation=None, *args, **kwargs):
  56. if not brightness is None and brightness>0:
  57. self.brightness = [max(1-brightness, 0), 1+brightness]
  58. if not contrast is None and contrast>0:
  59. self.contrast = [max(1-contrast, 0), 1+contrast]
  60. if not saturation is None and saturation>0:
  61. self.saturation = [max(1-saturation, 0), 1+saturation]
  62. def __call__(self, im_lb):
  63. im = im_lb['im']
  64. lb = im_lb['lb']
  65. r_brightness = random.uniform(self.brightness[0], self.brightness[1])
  66. r_contrast = random.uniform(self.contrast[0], self.contrast[1])
  67. r_saturation = random.uniform(self.saturation[0], self.saturation[1])
  68. im = ImageEnhance.Brightness(im).enhance(r_brightness)
  69. im = ImageEnhance.Contrast(im).enhance(r_contrast)
  70. im = ImageEnhance.Color(im).enhance(r_saturation)
  71. return dict(im = im,
  72. lb = lb,
  73. )
  74. class MultiScale(object):
  75. def __init__(self, scales):
  76. self.scales = scales
  77. def __call__(self, img):
  78. W, H = img.size
  79. sizes = [(int(W*ratio), int(H*ratio)) for ratio in self.scales]
  80. imgs = []
  81. [imgs.append(img.resize(size, Image.BILINEAR)) for size in sizes]
  82. return imgs
  83. class Compose(object):
  84. def __init__(self, do_list):
  85. self.do_list = do_list
  86. def __call__(self, im_lb):
  87. for comp in self.do_list:
  88. im_lb = comp(im_lb)
  89. return im_lb
  90. if __name__ == '__main__':
  91. flip = HorizontalFlip(p = 1)
  92. crop = RandomCrop((321, 321))
  93. rscales = RandomScale((0.75, 1.0, 1.5, 1.75, 2.0))
  94. img = Image.open('data/img.jpg')
  95. lb = Image.open('data/label.png')