refactors_needed

app.py

@@ -1,11 +1,9 @@
 import os
 import cv2
-from PIL import Image
-import pandas as pd
 import streamlit as st
 from ultralytics import YOLO
 from streamlit_image_comparison import image_comparison
-import numpy as np
+from utils import save_uploadedfile, apply_masks
 
 model = YOLO('models/last.pt')
 
@@ -23,47 +21,7 @@ url = "https://github.com/AkanimohOD19A/img-segmentation"
 link = f'<a href="{url}">This sample app was heavily based on code shared by Akan Daniel. Huge credits to him.</a>'
 st.markdown(link, unsafe_allow_html=True)
 
-def save_uploadedfile(uploadedfile):
-    with open(os.path.join(parent_media_path, "captured_picture.jpg"), "wb") as f:
-        im = Image.open(uploadedfile)
-        im.save(os.path.join(parent_media_path, "captured_picture.jpg"))
 
-def convert_to_jpg(uploaded_image):
-    im = Image.open(uploaded_image)
-    if im.mode in ("RGBA", "P"):
-        im = im.convert("RGB")
-    uploaded_image_path = os.path.join(parent_media_path, "uploaded_image.jpg")
-    im.save(uploaded_image_path)
-
-
-def overlay(image, mask, color, alpha, resize=None):
-    """Combines image and its segmentation mask into a single image.
-
-    Params:
-        image: Training image. np.ndarray,
-        mask: Segmentation mask. np.ndarray,
-        color: Color for segmentation mask rendering.  tuple[int, int, int] = (255, 0, 0)
-        alpha: Segmentation mask's transparency. float = 0.5,
-        resize: If provided, both image and its mask are resized before blending them together.
-        tuple[int, int] = (1024, 1024))
-
-    Returns:
-        image_combined: The combined image. np.ndarray
-
-    """
-    # color = color[::-1]
-    colored_mask = np.expand_dims(mask, 0).repeat(3, axis=0)
-    colored_mask = np.moveaxis(colored_mask, 0, -1)
-    masked = np.ma.MaskedArray(image, mask=colored_mask, fill_value=color)
-    image_overlay = masked.filled()
-
-    if resize is not None:
-        image = cv2.resize(image.transpose(1, 2, 0), resize)
-        image_overlay = cv2.resize(image_overlay.transpose(1, 2, 0), resize)
-
-    image_combined = cv2.addWeighted(image, 1 - alpha, image_overlay, alpha, 0)
-
-    return image_combined
 
 st.divider()
 
@@ -84,7 +42,7 @@ if APPLICATION_MODE == "Take Picture":
     st.sidebar.write(
         """
             A computer aided application that segments your input image, built on 
-            the powerful YOLOv8 object detection algorithm developed by *ultralytics*.
+            the powerful YOLOv8 instance segmentation algorithm developed by *ultralytics*.
 
             Simply take a captured_picture and it gets segmentated in real time.
         """
@@ -95,12 +53,9 @@ if APPLICATION_MODE == "Take Picture":
         st.sidebar.divider()
         st.sidebar.image(picture, caption="captured_picture")
         if st.button("Segment!"):
-            ## Function to save image
-
-            save_uploadedfile(picture)
-            st.sidebar.success("Saved File")
-            # captured_picture_img = os.path.join(parent_media_path, "captured_picture.jpg")
             img_file = os.path.join(parent_media_path, "captured_picture.jpg")
+            save_uploadedfile(picture, save_path=img_file)
+            st.sidebar.success("Saved File")
         st.write("Click on **Clear photo** to retake picture")
 
     st.divider()
@@ -115,45 +70,29 @@ elif APPLICATION_MODE == "Upload Picture":
         """
     )
     st.sidebar.divider()
-    # uploaded_file = st.sidebar.file_uploader("Upload your Image here", type=['png', 'jpeg', 'jpg'])
     uploaded_file = st.sidebar.file_uploader("Drop a JPG/PNG file", accept_multiple_files=False, type=['jpg', 'png'])
-    if uploaded_file is not None and uploaded_file.type != ".jpg":
-        convert_to_jpg(uploaded_file)
     if uploaded_file is not None:
+        img_file = os.path.join(parent_media_path, "uploaded_image.jpg")
+        save_uploadedfile(uploaded_file, save_path=os.path.join(parent_media_path, "uploaded_image.jpg"))
         file_details = {"FileName": uploaded_file.name, "FileType": uploaded_file.type}
-        new_file_name = "uploaded_image.jpg"
-        with open(os.path.join(parent_media_path, new_file_name), "wb") as f:
-            f.write(uploaded_file.getbuffer())
-        img_file = os.path.join(parent_media_path, new_file_name)
+
         st.sidebar.success("File saved successfully")
-        print(f"File saved successfully to {os.path.abspath(os.path.join(parent_media_path, new_file_name))}")
+        print(f"File saved successfully to {os.path.abspath(img_file)}")
     else:
         st.sidebar.write("You are using a placeholder image, Upload your Image (.jpg for now) to explore")
 
-# def make_segmentation(img_file):
 results = model(img_file)
 img = cv2.imread(img_file)
-names_list = []
 for result in results:
-    # boxes = result.boxes.cpu().numpy()
     # segmentation
-    if result.masks:
-        masks = result.masks.data.cpu().numpy()  # masks, (N, H, W)
-    else:
-        masks = []
-
+    masks = [] if not result.masks else result.masks.data.cpu().numpy()
     numCols = len(masks)
     if numCols > 0:
         cols = st.columns(numCols)
+        print(f"Number of instances found: {numCols}")
     else:
-        print(f"Number of Boxes found: {numCols}")
         st.warning("Unable to id Distinct items - Please retry with a clearer Image")
-    for mask in masks:
-        h, w, _ = img.shape
-        mask = cv2.resize(mask, (w, h))
-        img = overlay(img, mask, color=(0, 255, 0), alpha=0.3)
-#     st.image(rect)
-#     render image-comparison
+    img = apply_masks(img, masks)
 
     st.markdown('')
     st.markdown('##### Slider of Uploaded Image and Segments')
@@ -173,14 +112,6 @@ st.sidebar.divider()
 st.sidebar.markdown('')
 st.sidebar.markdown('#### Distribution of identified items')
 
-# Boolean to resize the dataframe, stored as a session state variable
-st.sidebar.checkbox("Use container width", value=False, key="use_container_width")
-if len(names_list) > 0:
-    df_x = pd.DataFrame(names_list)
-    summary_table = df_x[0].value_counts().rename_axis('unique_values').reset_index(name='counts')
-    st.sidebar.dataframe(summary_table, use_container_width=st.session_state.use_container_width)
-else:
-    st.sidebar.warning("Unable to id Distinct items - Please retry with a clearer Image")
 
 st.markdown('')
 st.markdown('')

utils.py

@@ -0,0 +1,47 @@
+import numpy as np
+from PIL import Image
+import cv2
+
+
+def save_uploadedfile(uploaded_image, save_path):
+    im = Image.open(uploaded_image)
+    if im.mode in ("RGBA", "P"):
+        im = im.convert("RGB")
+    im.save(save_path)
+
+
+def overlay(image, mask, color, alpha, resize=None):
+    """Combines image and its segmentation mask into a single image.
+
+    Params:
+        image: Training image. np.ndarray,
+        mask: Segmentation mask. np.ndarray,
+        color: Color for segmentation mask rendering.  tuple[int, int, int] = (255, 0, 0)
+        alpha: Segmentation mask's transparency. float = 0.5,
+        resize: If provided, both image and its mask are resized before blending them together.
+        tuple[int, int] = (1024, 1024))
+
+    Returns:
+        image_combined: The combined image. np.ndarray
+
+    """
+    colored_mask = np.expand_dims(mask, 0).repeat(3, axis=0)
+    colored_mask = np.moveaxis(colored_mask, 0, -1)
+    masked = np.ma.MaskedArray(image, mask=colored_mask, fill_value=color)
+    image_overlay = masked.filled()
+
+    if resize is not None:
+        image = cv2.resize(image.transpose(1, 2, 0), resize)
+        image_overlay = cv2.resize(image_overlay.transpose(1, 2, 0), resize)
+
+    image_combined = cv2.addWeighted(image, 1 - alpha, image_overlay, alpha, 0)
+
+    return image_combined
+
+
+def apply_masks(img, masks):
+    for mask in masks:
+        h, w, _ = img.shape
+        mask = cv2.resize(mask, (w, h))
+        img = overlay(img, mask, color=(0, 255, 0), alpha=0.3)
+    return img