Adding image crop and adjustments methods (color, contrast, brightness, sharpness) (#154)

* Adding image adjustments methods (color, contrast, brightness, sharpness)

* Typing

* Improving adjustment tests

* lint

Author: pappacena

landingai/pipeline/frameset.py

@@ -3,13 +3,13 @@
 
 import logging
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Union, cast
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, cast
 import warnings
 
 import cv2
 import imageio
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageEnhance
 from pydantic import BaseModel
 
 from landingai.common import ClassificationPrediction, OcrPrediction, Prediction
@@ -256,6 +256,59 @@ def downsize(
             self.image = self.image.resize((width, height))
         return self
 
+    def crop(self, bbox: Tuple[int, int, int, int]) -> "Frame":
+        """Crop the image based on the bounding box
+
+        Parameters
+        ----------
+        bbox: A tuple with the bounding box coordinates (xmin, ymin, xmax, ymax)
+        """
+        self.image = self.image.crop(bbox)
+        return self
+
+    def adjust_sharpness(self, factor: float) -> "Frame":
+        """Adjust the sharpness of the image
+
+        Parameters
+        ----------
+        factor: The enhancement factor
+        """
+        return self._apply_enhancement(ImageEnhance.Sharpness, factor)
+
+    def adjust_brightness(self, factor: float) -> "Frame":
+        """Adjust the brightness of the image
+
+        Parameters
+        ----------
+        factor: The enhancement factor
+        """
+        return self._apply_enhancement(ImageEnhance.Brightness, factor)
+
+    def adjust_contrast(self, factor: float) -> "Frame":
+        """Adjust the contrast of the image
+
+        Parameters
+        ----------
+        factor: The enhancement factor
+        """
+        return self._apply_enhancement(ImageEnhance.Contrast, factor)
+
+    def adjust_color(self, factor: float) -> "Frame":
+        """Adjust the color of the image
+
+        Parameters
+        ----------
+        factor: The enhancement factor
+        """
+        return self._apply_enhancement(ImageEnhance.Color, factor)
+
+    def _apply_enhancement(
+        self, enhancement: Type[ImageEnhance._Enhance], factor: float
+    ) -> "Frame":
+        enhancer = enhancement(self.image)  # type: ignore
+        self.image = enhancer.enhance(factor)
+        return self
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -380,6 +433,67 @@ def downsize(
             frame.downsize(width, height)
         return self
 
+    def crop(self, bbox: Tuple[int, int, int, int]) -> "FrameSet":
+        """Crop the images based on the bounding box
+
+        Parameters
+        ----------
+        bbox: A tuple with the bounding box coordinates (xmin, ymin, xmax, ymax)
+        """
+        for frame in self.frames:
+            frame.crop(bbox)
+        return self
+
+    def adjust_sharpness(self, factor: float) -> "FrameSet":
+        """Adjust the sharpness of the image
+
+        Parameters
+        ----------
+        factor: The enhancement factor
+        """
+        for f in self.frames:
+            f.adjust_sharpness(factor)
+        return self
+
+    def adjust_brightness(self, factor: float) -> "FrameSet":
+        """Adjust the brightness of the image
+
+        Parameters
+        ----------
+        factor: The enhancement factor
+        """
+        for f in self.frames:
+            f.adjust_brightness(factor)
+        return self
+
+    def adjust_contrast(self, factor: float) -> "FrameSet":
+        """Adjust the contrast of the image
+
+        Parameters
+        ----------
+        factor: The enhancement factor
+        """
+        for f in self.frames:
+            f.adjust_contrast(factor)
+        return self
+
+    def adjust_color(self, factor: float) -> "FrameSet":
+        """Adjust the color of the image
+
+        Parameters
+        ----------
+        factor: The enhancement factor
+        """
+        for f in self.frames:
+            f.adjust_color(factor)
+        return self
+
+    def copy(self, *args: Any, **kwargs: Any) -> "FrameSet":
+        """Returns a copy of this FrameSet, with all the frames copied"""
+        frameset = super().copy(*args, **kwargs)
+        frameset.frames = [frame.copy() for frame in self.frames]
+        return frameset
+
     def save_image(
         self,
         filename_prefix: str,

tests/data/images/cereal-tiny/brightness-1.5.jpeg

@@ -2,6 +2,7 @@
 from pathlib import Path
 from unittest import mock
 
+from numpy.testing import assert_array_equal
 from PIL import Image
 import pytest
 
@@ -10,16 +11,18 @@
 from landingai.pipeline.frameset import FrameSet, Frame, PredictionList
 
 
-def get_frameset() -> FrameSet:
-    return FrameSet.from_image("tests/data/images/cereal1.jpeg")
+def get_frameset(image_path: str = "tests/data/images/cereal1.jpeg") -> FrameSet:
+    return FrameSet.from_image(image_path)
 
 
-def get_frame() -> Frame:
-    return Frame.from_image("tests/data/images/cereal1.jpeg")
+def get_frame(image_path: str = "tests/data/images/cereal1.jpeg") -> Frame:
+    return Frame.from_image(image_path)
 
 
-def get_frameset_with_od_coffee_prediction() -> FrameSet:
-    frameset = get_frameset()
+def get_frameset_with_od_coffee_prediction(
+    image_path: str = "tests/data/images/cereal1.jpeg",
+) -> FrameSet:
+    frameset = get_frameset(image_path)
     frameset.frames[0].predictions = PredictionList(
         [
             ObjectDetectionPrediction(
@@ -34,8 +37,10 @@ def get_frameset_with_od_coffee_prediction() -> FrameSet:
     return frameset
 
 
-def get_frame_with_od_coffee_prediction() -> Frame:
-    frame = get_frame()
+def get_frame_with_od_coffee_prediction(
+    image_path: str = "tests/data/images/cereal1.jpeg",
+) -> Frame:
+    frame = get_frame(image_path)
     frame.predictions = PredictionList(
         [
             ObjectDetectionPrediction(
@@ -234,6 +239,106 @@ def test_frameset_downsize_bigger_than_original(frame_getter):
     assert get_image_from_frame_or_frameset(frame).size == (width, height)
 
 
+@pytest.mark.parametrize(
+    "frame_getter",
+    [
+        get_frameset_with_od_coffee_prediction,
+        get_frame_with_od_coffee_prediction,
+    ],
+)
+def test_crop(frame_getter):
+    frame = frame_getter()
+    image = get_image_from_frame_or_frameset(frame)
+    width, height = image.size
+    # Crop 5px from each side
+    frame.crop((5, 5, width - 5, height - 5))
+    assert get_image_from_frame_or_frameset(frame).size == (width - 10, height - 10)
+
+
+@pytest.mark.parametrize(
+    "frame_getter",
+    [
+        get_frameset_with_od_coffee_prediction,
+        get_frame_with_od_coffee_prediction,
+    ],
+)
+@pytest.mark.parametrize(
+    "enhancement",
+    [
+        ("adjust_sharpness", 1.0),
+        ("adjust_brightness", 1.0),
+        ("adjust_contrast", 1.0),
+        ("adjust_color", 1.0),
+    ],
+)
+def test_enhancements_no_op(enhancement, frame_getter):
+    """Checks each enhancement method with it's no-op factor, making sure the frame
+    will be exactly the same after the operation"""
+    enhancement_method, enhancement_factor = enhancement
+    frame = frame_getter()
+    original_frame = frame.copy()
+    getattr(frame, enhancement_method)(enhancement_factor)
+
+    original_image = get_image_from_frame_or_frameset(original_frame)
+    image = get_image_from_frame_or_frameset(frame)
+    assert original_image is not image
+    assert_array_equal(original_image, image)
+
+
+@pytest.mark.parametrize(
+    "frame_getter",
+    [
+        get_frameset_with_od_coffee_prediction,
+        get_frame_with_od_coffee_prediction,
+    ],
+)
+@pytest.mark.parametrize(
+    "enhancement",
+    [
+        ("adjust_sharpness", 1.5, "sharpness-1.5.jpeg"),
+        ("adjust_brightness", 1.5, "brightness-1.5.jpeg"),
+        ("adjust_contrast", 1.5, "contrast-1.5.jpeg"),
+        ("adjust_color", 1.5, "color-1.5.jpeg"),
+    ],
+)
+def test_enhancements(enhancement, frame_getter):
+    """Checks each enhancement method with it's factor, making sure the frame
+    wil be changed after the operation"""
+    enhancement_method, enhancement_factor, expected_img_file = enhancement
+    img_folder = "tests/data/images/cereal-tiny/"
+    frame = frame_getter(image_path=f"{img_folder}/original.jpeg")
+    original_frame = frame.copy()
+    getattr(frame, enhancement_method)(enhancement_factor)
+
+    original_image = get_image_from_frame_or_frameset(original_frame)
+    image = get_image_from_frame_or_frameset(frame)
+    assert original_image is not image
+
+    expected_content = Image.open(f"{img_folder}/{expected_img_file}")
+    assert_array_equal(image, expected_content)
+
+
+@pytest.mark.parametrize(
+    "frame_getter",
+    [
+        get_frameset_with_od_coffee_prediction,
+        get_frame_with_od_coffee_prediction,
+    ],
+)
+def test_copy_operation(frame_getter):
+    frame = frame_getter()
+    width, height = get_image_from_frame_or_frameset(frame).size
+    new_frame = frame.copy()
+    # Do any operation on the new frame, so we can test if the original frame is preserved
+    new_frame.crop((0, 0, 1, 1))
+
+    original_image = get_image_from_frame_or_frameset(frame)
+    new_image = get_image_from_frame_or_frameset(new_frame)
+    assert frame.predictions == new_frame.predictions
+    assert original_image is not new_image
+    assert original_image.size == (width, height)
+
+
 @pytest.mark.parametrize(
     "frame_getter",
     [