Pass PIL.Image.Image to BrightnessContrastDialog

labelmeai · Jun 14, 2020 · 83bc799 · 83bc799
1 parent a2ebb1f
commit 83bc799
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 18 deletions.
diff --git a/labelme/app.py b/labelme/app.py
@@ -1364,7 +1364,9 @@ def onNewBrightnessContrast(self, qimage):
 
     def brightnessContrast(self, value):
         dialog = BrightnessContrastDialog(
-            self.filename, self.onNewBrightnessContrast, parent=self
+            utils.img_data_to_pil(self.imageData),
+            self.onNewBrightnessContrast,
+            parent=self,
         )
         brightness, contrast = self.brightnessContrast_values.get(
             self.filename, (None, None)
@@ -1484,7 +1486,9 @@ def loadFile(self, filename=None):
                 )
         # set brightness constrast values
         dialog = BrightnessContrastDialog(
-            self.filename, self.onNewBrightnessContrast, parent=self
+            utils.img_data_to_pil(self.imageData),
+            self.onNewBrightnessContrast,
+            parent=self,
         )
         brightness, contrast = self.brightnessContrast_values.get(
             self.filename, (None, None)

diff --git a/labelme/utils/__init__.py b/labelme/utils/__init__.py
@@ -6,7 +6,9 @@
 from .image import img_arr_to_b64
 from .image import img_b64_to_arr
 from .image import img_data_to_arr
+from .image import img_data_to_pil
 from .image import img_data_to_png_data
+from .image import img_pil_to_data
 
 from .shape import labelme_shapes_to_label
 from .shape import masks_to_bboxes

diff --git a/labelme/utils/image.py b/labelme/utils/image.py
@@ -7,10 +7,16 @@
 import PIL.ImageOps
 
 
-def img_data_to_arr(img_data):
+def img_data_to_pil(img_data):
     f = io.BytesIO()
     f.write(img_data)
-    img_arr = np.array(PIL.Image.open(f))
+    img_pil = PIL.Image.open(f)
+    return img_pil
+
+
+def img_data_to_arr(img_data):
+    img_pil = img_data_to_pil(img_data)
+    img_arr = np.array(img_pil)
     return img_arr
 
 
@@ -20,6 +26,13 @@ def img_b64_to_arr(img_b64):
     return img_arr
 
 
+def img_pil_to_data(img_pil):
+    f = io.BytesIO()
+    img_pil.save(f, format="PNG")
+    img_data = f.getvalue()
+    return img_data
+
+
 def img_arr_to_b64(img_arr):
     img_pil = PIL.Image.fromarray(img_arr)
     f = io.BytesIO()

diff --git a/labelme/widgets/brightness_contrast_dialog.py b/labelme/widgets/brightness_contrast_dialog.py
@@ -1,13 +1,14 @@
+import PIL.Image
+import PIL.ImageEnhance
 from qtpy.QtCore import Qt
 from qtpy import QtGui
 from qtpy import QtWidgets
 
-from PIL import Image
-from PIL import ImageEnhance
+from .. import utils
 
 
 class BrightnessContrastDialog(QtWidgets.QDialog):
-    def __init__(self, filename, callback, parent=None):
+    def __init__(self, img, callback, parent=None):
         super(BrightnessContrastDialog, self).__init__(parent)
         self.setModal(True)
         self.setWindowTitle("Brightness/Contrast")
@@ -20,26 +21,25 @@ def __init__(self, filename, callback, parent=None):
         formLayout.addRow(self.tr("Contrast"), self.slider_contrast)
         self.setLayout(formLayout)
 
-        self.img = Image.open(filename).convert("RGBA")
+        assert isinstance(img, PIL.Image.Image)
+        self.img = img
         self.callback = callback
 
     def onNewValue(self, value):
-        brightness = self.slider_brightness.value() / 100.0
-        contrast = self.slider_contrast.value() / 100.0
+        brightness = self.slider_brightness.value() / 50.0
+        contrast = self.slider_contrast.value() / 50.0
 
         img = self.img
-        img = ImageEnhance.Brightness(img).enhance(brightness)
-        img = ImageEnhance.Contrast(img).enhance(contrast)
+        img = PIL.ImageEnhance.Brightness(img).enhance(brightness)
+        img = PIL.ImageEnhance.Contrast(img).enhance(contrast)
 
-        bytes = img.tobytes("raw", "RGBA")
-        qimage = QtGui.QImage(
-            bytes, img.size[0], img.size[1], QtGui.QImage.Format_RGB32
-        ).rgbSwapped()
+        img_data = utils.img_pil_to_data(img)
+        qimage = QtGui.QImage.fromData(img_data)
         self.callback(qimage)
 
     def _create_slider(self):
         slider = QtWidgets.QSlider(Qt.Horizontal)
-        slider.setRange(0, 300)
-        slider.setValue(100)
+        slider.setRange(0, 150)
+        slider.setValue(50)
         slider.valueChanged.connect(self.onNewValue)
         return slider