35
35
import copy
36
36
import logging
37
37
import cv2
38
- from PIL import Image , ImageDraw
38
+ from PIL import Image , ImageDraw , ImageEnhance
39
39
import pickle
40
40
import threading
41
41
MUTEX = threading .Lock ()
@@ -490,10 +490,10 @@ class RandomDistort(BaseOperator):
490
490
saturation (list): saturation settings. in [lower, upper, probability] format.
491
491
contrast (list): contrast settings. in [lower, upper, probability] format.
492
492
brightness (list): brightness settings. in [lower, upper, probability] format.
493
- random_apply (bool): whether to apply in random (yolo) or fixed (SSD)
494
- order .
495
- count (int ): the number of doing distrot
496
- random_channel (bool ): whether to swap channels randomly
493
+ random_apply (bool): whether to apply in random (yolo) or fixed (SSD) order.
494
+ count (int): the number of doing distrot .
495
+ random_channel (bool ): whether to swap channels randomly.
496
+ prob (float ): the probability of enhancing the sample.
497
497
"""
498
498
499
499
def __init__ (self ,
@@ -519,57 +519,41 @@ def apply_hue(self, img):
519
519
low , high , prob = self .hue
520
520
if np .random .uniform (0. , 1. ) < prob :
521
521
return img
522
-
523
- img = img .astype (np .float32 )
524
- # it works, but result differ from HSV version
525
522
delta = np .random .uniform (low , high )
526
- u = np .cos (delta * np .pi )
527
- w = np .sin (delta * np .pi )
528
- bt = np .array ([[1.0 , 0.0 , 0.0 ], [0.0 , u , - w ], [0.0 , w , u ]])
529
- tyiq = np .array ([[0.299 , 0.587 , 0.114 ], [0.596 , - 0.274 , - 0.321 ],
530
- [0.211 , - 0.523 , 0.311 ]])
531
- ityiq = np .array ([[1.0 , 0.956 , 0.621 ], [1.0 , - 0.272 , - 0.647 ],
532
- [1.0 , - 1.107 , 1.705 ]])
533
- t = np .dot (np .dot (ityiq , bt ), tyiq ).T
534
- img = np .dot (img , t )
523
+ img = np .array (img .convert ('HSV' ))
524
+ img [:, :, 0 ] = img [:, :, 0 ] + delta
525
+ img = Image .fromarray (img , mode = 'HSV' ).convert ('RGB' )
535
526
return img
536
527
537
528
def apply_saturation (self , img ):
538
529
low , high , prob = self .saturation
539
530
if np .random .uniform (0. , 1. ) < prob :
540
531
return img
541
532
delta = np .random .uniform (low , high )
542
- img = img .astype (np .float32 )
543
- # it works, but result differ from HSV version
544
- gray = img * np .array ([[[0.299 , 0.587 , 0.114 ]]], dtype = np .float32 )
545
- gray = gray .sum (axis = 2 , keepdims = True )
546
- gray *= (1.0 - delta )
547
- img *= delta
548
- img += gray
533
+ img = ImageEnhance .Color (img ).enhance (delta )
549
534
return img
550
535
551
536
def apply_contrast (self , img ):
552
537
low , high , prob = self .contrast
553
538
if np .random .uniform (0. , 1. ) < prob :
554
539
return img
555
540
delta = np .random .uniform (low , high )
556
- img = img .astype (np .float32 )
557
- img *= delta
541
+ img = ImageEnhance .Contrast (img ).enhance (delta )
558
542
return img
559
543
560
544
def apply_brightness (self , img ):
561
545
low , high , prob = self .brightness
562
546
if np .random .uniform (0. , 1. ) < prob :
563
547
return img
564
548
delta = np .random .uniform (low , high )
565
- img = img .astype (np .float32 )
566
- img += delta
549
+ img = ImageEnhance .Brightness (img ).enhance (delta )
567
550
return img
568
551
569
552
def apply (self , sample , context = None ):
570
553
if random .random () > self .prob :
571
554
return sample
572
555
img = sample ['image' ]
556
+ img = Image .fromarray (img .astype (np .uint8 ))
573
557
if self .random_apply :
574
558
functions = [
575
559
self .apply_brightness , self .apply_contrast ,
@@ -578,21 +562,20 @@ def apply(self, sample, context=None):
578
562
distortions = np .random .permutation (functions )[:self .count ]
579
563
for func in distortions :
580
564
img = func (img )
565
+ img = np .asarray (img ).astype (np .float32 )
581
566
sample ['image' ] = img
582
567
return sample
583
568
584
569
img = self .apply_brightness (img )
585
570
mode = np .random .randint (0 , 2 )
586
-
587
571
if mode :
588
572
img = self .apply_contrast (img )
589
-
590
573
img = self .apply_saturation (img )
591
574
img = self .apply_hue (img )
592
-
593
575
if not mode :
594
576
img = self .apply_contrast (img )
595
577
578
+ img = np .asarray (img ).astype (np .float32 )
596
579
if self .random_channel :
597
580
if np .random .randint (0 , 2 ):
598
581
img = img [..., np .random .permutation (3 )]
0 commit comments