Merge pull request #36 from Flippchen/dev

Dev

Author: Flippchen

.github/workflows/build_binary.yaml

@@ -15,10 +15,10 @@ jobs:
     - name: Check out repository
       uses: actions/checkout@v2
 
-    - name: Set up Python 3.9
+    - name: Set up Python 3.10
       uses: actions/setup-python@v2
       with:
-        python-version: 3.9
+        python-version: '3.10'
 
     - name: Install dependencies
       run: |
@@ -43,10 +43,10 @@ jobs:
     - name: Check out repository
       uses: actions/checkout@v2
 
-    - name: Set up Python 3.9
+    - name: Set up Python 3.10
       uses: actions/setup-python@v2
       with:
-        python-version: 3.9
+        python-version: '3.10'
 
     - name: Install dependencies
       run: |
@@ -73,10 +73,10 @@ jobs:
       - name: Check out repository
         uses: actions/checkout@v2
 
-      - name: Set up Python 3.9
+      - name: Set up Python 3.10
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9
+          python-version: '3.10'
 
       - name: Install dependencies
         run: |

.github/workflows/python.yaml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ "3.9", "3.10" ]
+        python-version: ['3.10']
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}

README.md

@@ -32,7 +32,7 @@ python web_app/main.py
 or download the [executable](https://github.com/Flippchen/PorscheInsight-CarClassification-AI/actions).
 
 ### Demo
-<img alt="Screenshot of the Web UI" src="assets/web_app/demo_web_ui.gif"  height="500">
+<img alt="Screenshot of the Web UI" src="assets/web_app/demo_web_ui_mask.gif"  height="500">
 
 ### Architecture
 The Web UI employs a two-step process involving two models. Initially, the pre_filter model determines if an image contains a Porsche. If a Porsche is detected, the image proceeds to the second model, which classifies the car according to the user's input.
@@ -46,6 +46,7 @@ To see the architecture of the Local App UI, click the arrow below.
 </details>
 
 ### ToDos
+- [ ] Build an ensemble model
 - [ ] Switch to Google Cloud function/use S3 bucket/compress image
 - [ ] Improve pre_filter model
 - [ ] Retrain car type model with efficientnet model
@@ -72,6 +73,7 @@ To see the architecture of the Local App UI, click the arrow below.
 - [x] Implement new architecture to online version
 - [x] Add release 1.0.0
 - [x] Add docker support (AWS Lambda)
+- [x] Display the mask of the predicted car
 </details>
 
 

assets/web_app/demo_web_ui.gif

@@ -7,6 +7,7 @@ keras~=2.10.0
 # Inference
 rembg==2.0.32
 onnxruntime==1.14.1
+scipy~=1.9.3
 # Export to onnx
 tf2onnx==1.14.0
 # Explaining the model

assets/web_app/demo_web_ui_mask.gif

@@ -16,9 +16,11 @@ def test_get_session():
 def test_replace_background():
     # Test using a sample image
     im = Image.new('RGB', (100, 100), 'WHITE')
-    new_im = replace_background(im)
+    new_im, mask = replace_background(im)
     assert new_im is not None
     assert isinstance(new_im, Image.Image)
+    assert mask is not None
+    assert isinstance(mask, Image.Image)
 
 
 def test_get_bounding_box():
@@ -55,9 +57,11 @@ def test_load_and_remove_bg(tmpdir):
     test_image = Image.new('RGB', (100, 100), 'WHITE')
     test_image_path = os.path.join(tmpdir, 'test_image.jpg')
     test_image.save(test_image_path)
-    loaded_im = load_and_remove_bg(test_image_path, (50, 50))
+    loaded_im, mask = load_and_remove_bg(test_image_path, (50, 50))
     assert loaded_im is not None
     assert loaded_im.size == (50, 50)
+    assert mask is not None
+    assert mask.size == (50, 50)
 
 
 def test_remove_bg_from_all_images(tmpdir):

requirements.txt

@@ -1,22 +1,28 @@
 import io
 import os
-from PIL import Image, ImageOps, ImageFile
+
+import numpy as np
+from scipy.ndimage import maximum_filter
+from PIL import Image, ImageOps, ImageFile, ImageFilter, ImageChops
 from rembg import remove, new_session
 from PIL.Image import Image as PILImage
+from typing import Tuple
 from functools import cache
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
+
 @cache
 def get_session():
     return new_session("u2net")
 
 
-def replace_background(im: PILImage, post_process_mask=False, session=None, size: tuple = None) -> PILImage:
+def replace_background(im: PILImage, post_process_mask=True, session=None, size: tuple = None) -> Tuple[PILImage, PILImage]:
     size = size or (300, 300)
     # if not isinstance(im, PILImage):
     #   im = Image.open(io.BytesIO(im))
     session = session or get_session()
     im = remove(im, post_process_mask=post_process_mask, session=session)
+    mask = im.copy()
     im = resize_cutout(im, size)
 
     new_im = Image.new('RGBA', im.size, 'BLACK')
@@ -28,7 +34,7 @@ def replace_background(im: PILImage, post_process_mask=False, session=None, size
     image = Image.open(io.BytesIO(im_bytes))
     image = image.convert('RGB')
 
-    return image
+    return image, mask
 
 
 def get_bounding_box(im: PILImage) -> tuple:
@@ -104,19 +110,62 @@ def fix_image(image):
     return image
 
 
+def convert_mask(mask, color=(29, 132, 181), border_color=(219, 84, 97), border_fraction=0.03):
+    # Convert the image to RGBA if it is not already
+    if mask.mode != 'RGBA':
+        mask = mask.convert('RGBA')
+
+    border_size = int(min(mask.size) * border_fraction)
+
+    if border_size % 2 == 0:
+        border_size += 1
+
+    # Create a copy of the mask and expand it to create the border
+    mask_np = np.array(mask)
+    border_mask_np = maximum_filter(mask_np, size=border_size)
+
+    # Convert back to PIL image
+    border_mask = Image.fromarray(border_mask_np)
+
+    # Create the border by subtracting the original mask from the expanded mask
+    border = ImageChops.difference(border_mask, mask)
+
+    # Convert the mask and the border to the desired colors
+    mask_np = np.array(mask)
+    border_np = np.array(border)
+
+    # Mask the areas where alpha channel is not zero
+    mask_area = mask_np[..., 3] != 0
+    border_area = border_np[..., 3] != 0
+
+    # Replace RGB channels with desired color while keeping alpha channel the same
+    mask_np[mask_area, :3] = color
+    mask_np[mask_area, 3] = mask_np[mask_area, 3] // 4
+    border_np[border_area, :3] = border_color
+
+    # Convert back to PIL images
+    mask = Image.fromarray(mask_np)
+    border = Image.fromarray(border_np)
+
+    # Combine the mask and the border
+    mask_with_border = Image.alpha_composite(mask, border)
+
+    return mask_with_border
+
+
 def load_and_remove_bg(path, size):
     image = Image.open(path)
     # image = resize_image(image, size)
     image = resize_and_pad_image(image, size)
-    image = replace_background(image, size=size)
+    image, mask = replace_background(image, size=size)
 
-    return image
+    return image, mask
 
 
 def remove_bg_from_all_images(folder: str):
     for image in os.listdir(f'{folder}'):
         print("Removing background from", image)
-        img = load_and_remove_bg(f"{folder}/{image}", (300, 300))
+        img, mask = load_and_remove_bg(f"{folder}/{image}", (300, 300))
         img.save(f"{folder}/{image}")
 
 

tests/test_prepare_images.py

@@ -12,8 +12,9 @@
 import base64
 from io import BytesIO
 from PIL import Image
+
 from utilities.class_names import get_classes_for_model
-from utilities.prepare_images import replace_background, resize_and_pad_image, fix_image
+from utilities.prepare_images import replace_background, resize_and_pad_image, fix_image, convert_mask
 import pooch
 from rembg import new_session
 
@@ -61,14 +62,24 @@ def load_model(model_name: str) -> ort.InferenceSession:
     return ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
 
 
-def prepare_image(image_data: Image, target_size: Tuple, remove_background: bool) -> np.ndarray:
-    if remove_background:
-        image = replace_background(image_data, session=session)
+def prepare_image(image_data: Image, target_size: Tuple, remove_background: bool, show_mask: bool) -> Tuple[np.ndarray, Image.Image]:
+    if remove_background and show_mask:
+        image, mask = replace_background(image_data, session=session)
+    elif remove_background:
+        image, _ = replace_background(image_data, session=session)
+        mask = None
     else:
         image = resize_and_pad_image(image_data, target_size)
+        mask = None
+
     img_array = np.array(image).astype('float32')
     img_array = np.expand_dims(img_array, 0)
-    return img_array
+
+    if mask is None:
+        mask = Image.fromarray(np.zeros((1, 1), dtype=np.uint8))
+        mask = mask.convert("RGBA")
+
+    return img_array, mask
 
 
 def get_top_3_predictions(prediction: np.ndarray, model_name: str) -> List[Tuple[str, float]]:
@@ -88,35 +99,53 @@ def get_pre_filter_prediction(image_data: np.ndarray, model_name: str):
 
 
 @eel.expose
-def classify_image(image_data: str, model_name: str) -> List[Tuple[str, float]]:
-    # Load model if not loaded yet
+def classify_image(image_data: str, model_name: str, show_mask: str = "no") -> tuple[list[tuple[str, float]], str] | list[list[tuple[str, float]]]:
+    # Loading the model if it's not already loaded
     if models[model_name] is None:
         models[model_name] = load_model(model_name)
 
-    # Decode image and open it
+    # Decoding the base64 image data and opening the image
     image_data = base64.b64decode(image_data)
     image = Image.open(BytesIO(image_data))
+    show_mask = True if show_mask == "yes" else False
 
-    # Fix image orientation and color mode if needed
+    # Correcting image orientation and color mode if necessary
     image = fix_image(image)
-    # Get correct input size for model
+
+    # Retrieving the required input size for the specified model
     input_size = models[model_name].get_inputs()[0].shape[1:3]
 
-    # Prepare image for filtering and predict
-    # FIXME: It is a workaround to remove the bg in first place/ Why is the prediction with black background better?
-    filter_image = prepare_image(image, input_size, remove_background=True)
+    # Preparing image for processing and prediction
+    # FIXME: Currently, the background is removed prior to prediction. This is a workaround,
+    # as predictions seem to be better with a black background.
+    filter_image, mask = prepare_image(image, input_size, remove_background=True, show_mask=show_mask)
+
+    # Converting the mask for processing
+    mask = convert_mask(mask)
+    buffer = io.BytesIO()
+    mask.save(buffer, format="PNG")
+    mask_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+    # Getting initial predictions before applying filters
     filter_predictions = get_pre_filter_prediction(filter_image, "pre_filter")
-    # If the pre_filter predicts porsche or other_car_brand, predict the correct model
+
+    # If the initial prediction is 'porsche' or other car brands, run the specified model's prediction
     if filter_predictions[0][0] == "porsche":
-        # prepared_image = prepare_image(image, input_size, remove_background=True)
         input_name = models[model_name].get_inputs()[0].name
         prediction = models[model_name].run(None, {input_name: filter_image})
-        # Get top 3 predictions
+
+        # Retrieving the top 3 predictions
         top_3 = get_top_3_predictions(prediction[0], model_name)
 
-        return top_3
+        if show_mask:
+            return top_3, mask_base64
+        else:
+            return [top_3]
 
-    return filter_predictions
+    # Returning the initial predictions and an optional mask if show_mask is set to 'yes'
+    if show_mask:
+        return filter_predictions, mask_base64
+    return [filter_predictions]
 
 
 eel.init("web")

utilities/prepare_images.py

@@ -45,3 +45,8 @@ input[type="password"].dark-mode,
   color: #ffffff;
   border: 1px solid #6c757d;
 }
+#show-mask.dark-mode {
+  background-color: #3c4049;
+  color: #ffffff;
+  border: 1px solid #6c757d;
+}