Upload 3 files

app.py

@@ -2,13 +2,15 @@ import streamlit as st
 from rembg import remove
 from PIL import ImageOps, ImageEnhance, Image
 from streamlit_option_menu import option_menu
-from markup import real_estate_app, real_estate_app_hf, sliders_intro, perspective_intro, manual_bg_intro
+from markup import real_estate_app, real_estate_app_hf, sliders_intro, perspective_intro, manual_bg_intro, segement_intro
 from perspective_correction import perspective_correction, perspective_correction2
 from streamlit_drawable_canvas import st_canvas
 import tempfile
-
+from ultralytics import YOLO
 import numpy as np
 import cv2
+import gdown
+import os
 
 
 def tab1():
@@ -120,7 +122,66 @@ def remove_background(image, points):
     
     return result
 
+
 def tab3():
+    model_file = 'bestslab-seg.onnx'
+    model_url = 'https://drive.google.com/uc?id=1---iqs2llLrgDbzr_S1nzkKUr3sJ_ru3'
+
+    if not os.path.exists(model_file):
+        gdown.download(model_url, model_file, quiet=False)
+    
+    model = YOLO(model_file)
+
+    st.header("Background Removal with instance Segmentaion")
+    st.markdown(segement_intro(), unsafe_allow_html=True)
+
+    uploaded_file = st.file_uploader('Choose an image', type=['jpg', 'jpeg', 'png'])
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+        col1, col2 = st.columns([2,1])
+        with col2:
+
+            iou_threshold = st.slider('IoU Threshold', min_value=0.0, max_value=1.0, value=0.7)
+            conf_threshold = st.slider('Confidence Threshold', min_value=0.0, max_value=1.0, value=0.65)
+            show_labels = st.checkbox('Show Labels', value=False)
+            show_conf = st.checkbox('Show Confidence Scores', value=False)
+            boxes = st.checkbox('Show Boxes', value=True)
+
+        with col1:
+            st.image(image, caption='Input Image', use_column_width=True)
+
+            if st.button('Apply and Predict'):
+
+                results = model(
+                    image_cv,
+                    iou=iou_threshold,
+                    conf=conf_threshold,
+                    show_labels=show_labels,
+                    show_conf=show_conf,
+                    boxes=boxes,
+                )
+
+                masks = results[0].masks
+                mask_image = np.zeros((image_cv.shape[0], image_cv.shape[1], 4), dtype=np.uint8)
+
+                annotated_frame = results[0].plot()
+                annotated_image = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
+
+                for segment in masks.xy:
+                    segment = np.array(segment, dtype=np.int32)
+                    segment = segment.reshape((-1, 1, 2))
+                    cv2.fillPoly(mask_image, [segment], (255, 255, 255, 255))
+
+                alpha_channel = mask_image[:, :, 3]
+                image_rgba = np.concatenate((image_cv, np.expand_dims(alpha_channel, axis=2)), axis=2)
+                masked_image = image_rgba * (mask_image / 255)
+                masked_pil = Image.fromarray(masked_image.astype(np.uint8), 'RGBA')
+
+                st.image([annotated_image, masked_pil], caption=['Detections', 'Masked Image'], use_column_width=True)    
+
+def tab4():
     st.header("Manual Background Removal")
     st.markdown(manual_bg_intro(), unsafe_allow_html=True)
     uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
@@ -139,6 +200,7 @@ def tab3():
         with col2:
             drawing_mode = "point"
             stroke_width = st.slider("Stroke width: ", 1, 25, 3)
+            #point_display_radius = st.slider("Point display radius: ", 1, 25, 3)
             realtime_update = st.checkbox("Update in realtime", True)
 
         with col1:
@@ -151,10 +213,13 @@ def tab3():
                 height=image_height,
                 width=image_width,
                 drawing_mode=drawing_mode,
+                #point_display_radius=point_display_radius if drawing_mode == 'point' else 0,
                 key="canvas",
             )
 
         if st.button("Remove Background"):
+            st.subheader("This Feature is broken in this hosting service, please DM for Private link")
+            """
             if canvas_result.json_data is not None:
                 points = []
                 for obj in canvas_result.json_data["objects"]:
@@ -171,10 +236,10 @@ def tab3():
                 transparent_bg_result = result_image.convert("RGBA")
                 file_path = "background_removed.png"
                 transparent_bg_result.save(file_path, format="PNG")
-                st.image(transparent_bg_result, caption="Background Removed Image")   
+                st.image(transparent_bg_result, caption="Background Removed Image")   """
 
 
-def tab4():
+def tab5():
     st.header("Image Perspective Correction")
     st.write("Upload a transparent PNG image which you have removed the background using the previous tab.")
     st.markdown(perspective_intro(),unsafe_allow_html=True)
@@ -220,7 +285,7 @@ def tab4():
                 
 def main():
     st.set_page_config(page_title="Background Removal Demo", page_icon=":memo:", layout="wide")
-    tabs = ["Intro", "AI Background Removal", "Manual Background Removal", "Perspective Correction"]
+    tabs = ["Intro", "AI Background Removal", "Background Removal with Segmentaion", "Manual Background Removal", "Perspective Correction"]
 
     with st.sidebar:
 
@@ -229,12 +294,13 @@ def main():
     tab_functions = {
     "Intro": tab1,
     "AI Background Removal": tab2,
-    "Manual Background Removal": tab3,
-    "Perspective Correction": tab4,
+    "Background Removal with Segmentaion": tab3,
+    "Manual Background Removal": tab4,
+    "Perspective Correction": tab5,
     }
 
     if current_tab in tab_functions:
         tab_functions[current_tab]()
 
 if __name__ == "__main__":
-    main()
+    main()

bestslab-seg.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c61cc93ff7fc7477e12072c8ec4d1c96a7e302a33fb4a7cc8366b516c449d77a
+size 183900291

markup.py

@@ -68,4 +68,11 @@ def perspective_intro():
 def manual_bg_intro():
     return """
     Click on the canvas and select four or more points, starting from any one corner of the slab and proceeding in order in any direction. The background will be removed using the selected points.
+    """
+
+def segement_intro():
+    return """
+    Uses a custom trained model on multiple variations of the provided images. Accucary will increase if more images are provided and model is fine tuned. 
+    fine tuning the model for new images is a straight forward process, necessary documentaion and guidlines will be provided.
+    Adjust the confidence score and IoU sliders to refine the selected area.
     """