Upload 5 files

app.py

@@ -0,0 +1,240 @@
+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 perspective_correction import perspective_correction, perspective_correction2
+from streamlit_drawable_canvas import st_canvas
+import tempfile
+
+import numpy as np
+import cv2
+
+
+def tab1():
+    st.header("Image Background Remover")  
+    col1, col2 = st.columns([1, 2])
+    with col1:
+        st.image("image.jpg", use_column_width=True)
+    with col2:
+        st.markdown(real_estate_app(), unsafe_allow_html=True)
+    st.markdown(real_estate_app_hf(),unsafe_allow_html=True) 
+
+
+    github_link = '[<img src="https://badgen.net/badge/icon/github?icon=github&label">](https://github.com/ethanrom)'
+    huggingface_link = '[<img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue">](https://huggingface.co/ethanrom)'
+
+    st.write(github_link + '&nbsp;&nbsp;&nbsp;' + huggingface_link, unsafe_allow_html=True)
+
+def tab2():
+    st.header("Image Background Remover")
+    st.markdown(sliders_intro(), unsafe_allow_html=True)
+
+    upload_option = st.radio("Upload Option", ("Single Image", "Multiple Images"))
+
+    if upload_option == "Single Image":
+        uploaded_images = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"], accept_multiple_files=False)
+        images = [Image.open(uploaded_images)] if uploaded_images else []
+    else:
+        uploaded_images = st.file_uploader("Upload multiple images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
+        images = [Image.open(image) for image in uploaded_images] if uploaded_images else []
+
+    if images:
+        col1, col2 = st.columns([2, 1])
+
+        with col1:
+            st.image(images[0], caption="Original Image", use_column_width=True)
+            image = preprocess_image_1(images[0])
+
+        with col2:
+            st.subheader("RGB Adjustments")
+            with st.expander("Expand"):
+                r_min, r_max = st.slider("Red", min_value=0, max_value=255, value=(0, 255), step=1)
+                g_min, g_max = st.slider("Green", min_value=0, max_value=255, value=(0, 255), step=1)
+                b_min, b_max = st.slider("Blue", min_value=0, max_value=255, value=(0, 255), step=1)
+
+                adjusted_image = adjust_rgb(image, r_min, r_max, g_min, g_max, b_min, b_max)
+                st.image(adjusted_image, caption="Adjusted Image", use_column_width=True)
+
+            st.subheader("Curves Adjustment")
+            with st.expander("Expand"):
+                r_curve = st.slider("Red Curve", min_value=0.0, max_value=1.0, value=1.0, step=0.05)
+                g_curve = st.slider("Green Curve", min_value=0.0, max_value=1.0, value=1.0, step=0.05)
+                b_curve = st.slider("Blue Curve", min_value=0.0, max_value=1.0, value=1.0, step=0.05)
+
+                adjusted_image = adjust_curves(adjusted_image, r_curve, g_curve, b_curve)
+                st.image(adjusted_image, caption="Adjusted Image", use_column_width=True)
+
+            st.subheader("Masking")
+            with st.expander("Expand"):
+                threshold = st.slider("Threshold", min_value=0, max_value=255, value=128, step=1)
+
+                adjusted_image = apply_masking(adjusted_image, threshold)
+                st.image(adjusted_image, caption="Adjusted Image", use_column_width=True)
+
+        with col1:
+            if st.button("Remove Background"):
+                with st.spinner("Removing background..."):
+                    output_images = []
+                    for image in images:
+                        processed_image = preprocess_image_1(image)
+                        adjusted_image = adjust_rgb(processed_image, r_min, r_max, g_min, g_max, b_min, b_max)
+                        adjusted_image = adjust_curves(adjusted_image, r_curve, g_curve, b_curve)
+                        adjusted_image = apply_masking(adjusted_image, threshold)
+                        output_images.append(remove(adjusted_image))
+                    
+                    with st.expander("Background Removed Images"):
+                        for i in range(len(output_images)):
+                            st.image(output_images[i], caption=f"Background Removed Image {i + 1}", use_column_width=True)
+                        
+def preprocess_image_1(image):
+    if image.mode != "RGBA":
+        image = image.convert("RGBA")
+    return image
+
+def adjust_rgb(image, r_min, r_max, g_min, g_max, b_min, b_max):
+    r, g, b, a = image.split()
+    r = ImageOps.autocontrast(r.point(lambda p: int(p * (r_max - r_min) / 255 + r_min)))
+    g = ImageOps.autocontrast(g.point(lambda p: int(p * (g_max - g_min) / 255 + g_min)))
+    b = ImageOps.autocontrast(b.point(lambda p: int(p * (b_max - b_min) / 255 + b_min)))
+    return Image.merge("RGBA", (r, g, b, a))
+
+def adjust_curves(image, r_curve, g_curve, b_curve):
+    r, g, b, a = image.split()
+    enhancer_r = ImageEnhance.Brightness(r).enhance(r_curve)
+    enhancer_g = ImageEnhance.Brightness(g).enhance(g_curve)
+    enhancer_b = ImageEnhance.Brightness(b).enhance(b_curve)
+    return Image.merge("RGBA", (enhancer_r, enhancer_g, enhancer_b, a))
+
+def apply_masking(image, threshold):
+    r, g, b, a = image.split()
+    mask = a.point(lambda p: 255 if p > threshold else 0)
+    return Image.merge("RGBA", (r, g, b, mask))
+
+def remove_background(image, points):
+    mask = np.zeros(image.shape[:2], dtype=np.uint8)
+    points = np.array(points, dtype=np.int32)
+    points = points.reshape((-1, 1, 2))
+    cv2.fillPoly(mask, [points], (255))
+    result = np.dstack([image, mask])
+    
+    return result
+
+def tab3():
+    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"])
+    
+    if uploaded_file is not None:
+
+        col1, col2 = st.columns([2,1])
+        with col1:
+            image = Image.open(uploaded_file)
+            max_image_size = 700
+            if max(image.size) > max_image_size:
+                image.thumbnail((max_image_size, max_image_size), Image.LANCZOS)  # Updated resampling filter
+            st.image(image, caption="Original Image")
+            image_width, image_height = image.size
+
+        with col2:
+            drawing_mode = "point"
+            stroke_width = st.slider("Stroke width: ", 1, 25, 3)
+            realtime_update = st.checkbox("Update in realtime", True)
+
+        with col1:
+            st.subheader("Select Points on the Canvas")
+            canvas_result = st_canvas(
+                fill_color="rgba(255, 165, 0, 0.3)",
+                stroke_width=stroke_width,
+                background_image=image,
+                update_streamlit=realtime_update,
+                height=image_height,
+                width=image_width,
+                drawing_mode=drawing_mode,
+                key="canvas",
+            )
+
+        if st.button("Remove Background"):
+            if canvas_result.json_data is not None:
+                points = []
+                for obj in canvas_result.json_data["objects"]:
+                    if "type" in obj and obj["type"] == "circle":
+                        x = obj["left"]
+                        y = obj["top"]
+                        points.append((x, y))
+
+                img_array = np.array(image)
+                result = remove_background(img_array, points)
+
+                result_image = Image.fromarray(result)
+
+                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")   
+
+
+def tab4():
+    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)
+
+    uploaded_file = st.file_uploader("Choose a PNG image", type="png")
+
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+    
+        col1, col2 = st.columns([2,1])
+        with col1:
+            st.image(image, caption="Original Image", use_column_width=True)
+            image_np = np.array(image)
+
+        with col2:
+            correction_method = st.selectbox("Correction Method", ["Four-Point Perspective Correction", "Convex Hull Homography Perspective Correction"])
+
+            if correction_method == "Four-Point Perspective Correction":
+                threshold_value = st.slider("Threshold Value", min_value=1, max_value=255, value=100)
+                min_line_length = st.slider("Minimum Line Length", min_value=1, max_value=500, value=100)
+                max_line_gap = st.slider("Maximum Line Gap", min_value=1, max_value=100, value=10)
+            elif correction_method == "Convex Hull Homography Perspective Correction":
+                threshold_value = st.slider("Threshold Value", min_value=1, max_value=255, value=100)
+                min_line_length = st.slider("Minimum Line Length", min_value=1, max_value=500, value=100)
+                max_line_gap = st.slider("Maximum Line Gap", min_value=1, max_value=100, value=10)
+            else:
+                st.write("Invalid correction method selected.")
+                return
+
+        with col1:
+            if st.button("Correct Perspective"):
+                with st.spinner("Correcting Perspective..."):
+                    if uploaded_file is not None:
+                        if correction_method == "Four-Point Perspective Correction":
+                            corrected_image = perspective_correction(image_np, threshold_value, min_line_length, max_line_gap)
+                        elif correction_method == "Convex Hull Homography Perspective Correction":
+                            corrected_image = perspective_correction2(image_np, threshold_value, min_line_length, max_line_gap)
+                        else:
+                            st.write("Invalid correction method selected.")
+                            return
+
+                        st.image(corrected_image, caption="Corrected Image", use_column_width=True)
+                
+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"]
+
+    with st.sidebar:
+
+        current_tab = option_menu("Select a Tab", tabs, menu_icon="cast")
+
+    tab_functions = {
+    "Intro": tab1,
+    "AI Background Removal": tab2,
+    "Manual Background Removal": tab3,
+    "Perspective Correction": tab4,
+    }
+
+    if current_tab in tab_functions:
+        tab_functions[current_tab]()
+
+if __name__ == "__main__":
+    main()

markup.py

@@ -0,0 +1,71 @@
+def real_estate_app():
+    return """
+    <h3 style='text-align: center;'>Introduction</h3>
+            
+    <p>This app allows you to upload an image and remove its background using a custom trained model.</p>
+
+    <h4>Information:</h4>
+    <ul>
+        <li><b>Very low requirements:</b> This demo is running on a free single-core CPU. The app will work on any hardware. It may take a few minutes for the first run, but afterwards, it will perform faster.</li>
+        <li><b>Backend:</b> Python backend with an optimized ONNX model. No third party APIs or paid services used. Everything used is opensource</li>
+        <li><b>Frontend:</b> The demo's frontend is written in Python Streamlit. Unfortunately, I cannot write a JavaScript frontend as I am a backend engineer. However, I can refer you to a colleague if you prefer a better-looking UI.</li>
+        <li><b>Perspective correction:</b> Perspective correction is done via OpenCV.</li>
+    </ul>
+    </div>          
+    """
+
+def real_estate_app_hf():
+    return """
+    <div style='text-align: left;'>
+    <h3 style='text-align: center;'>About this background removal demo</h3>
+    <p>In this demo, the two requested features have been split into two separate tabs for better individual evaluation. They can be combined into one in the final submission.</p>
+    <br>
+    <h4>How to use:</h4>
+    <ul>
+    <li><b>Remove Background:</b> Use this tab to remove the background of the images. Right-click and save the image once done.</li>
+    <li><b>Correct Perspective:</b> Use this tab to check the perspective correction option. Please make sure to use an image with the background already removed.</li>
+    </ul>
+    <br>
+    <p><b>Update 1</b> RGB adjustment sliders for image preprocessing</p>
+    <p><b>Update 2</b> Addional options for perspective correction</p>
+    <p><b>Update 3</b> Multiple image support</p>
+    <p><b>Update 4</b> Manual background removal</p>
+    </div>
+    """
+
+def sliders_intro():
+    return """
+    <p>Newly added sliders which will appear after an image is uploaded serve as interactive tools to adjust various aspects of the uploaded image. Each slider corresponds to a specific color channel (red, green, blue) or a curve adjustment. 
+    By using these sliders, users can fine-tune the color levels and apply curve modifications to achieve the desired visual effect.</p>
+    <p>For the RGB Adjustments section, users can use the Red, Green, and Blue sliders to set the minimum and maximum values for each color channel. 
+    By adjusting these values, users can enhance or reduce the intensity of each color channel individually, allowing for precise color adjustments.</p>
+    <p>In the Curves Adjustment section, users can utilize the Red Curve, Green Curve, and Blue Curve sliders to control the brightness of the respective color channels. 
+    By moving these sliders, users can create custom curves, influencing the overall tone and contrast of the image.</p>
+    <p>The Masking section offers the Threshold slider, which determines the cutoff point for the transparency mask. 
+    Users can adjust the threshold value to define the boundary between foreground and background elements. This feature enables users to isolate objects by selectively applying transparency to specific areas of the image.</p>
+    """
+
+def perspective_intro():
+    return """
+    there are two different perspective correction methods you can chose from, the difference is how they determine the transformation matrix used for warping the image.
+
+    In the Four-Point Perspective Correction, the method uses a four-point perspective transform. It first detects lines in the image using the HoughLinesP function, and then calculates the endpoints of these lines. 
+    If enough endpoints are found (at least 4), a convex hull is created based on these endpoints. From the convex hull, a four-sided polygon is obtained, representing the region of interest. 
+    The width and height of this polygon are determined, and a destination set of points is defined to which the polygon will be mapped. 
+    Finally, a perspective transformation matrix is computed using getPerspectiveTransform function, and the image is warped accordingly.
+
+    In the Convex Hull Homography Perspective Correction, a similar process is followed, but instead of using a four-point perspective transform, it uses a homography transform. 
+    After obtaining the endpoints, a convex hull is created, and a four-sided polygon is extracted from it. The width and height of this polygon are calculated, and a destination set of points is defined. 
+    But instead of using getPerspectiveTransform, the findHomography function is used to compute the homography matrix. 
+    This matrix defines the transformation between the source polygon and the destination polygon, and the image is warped using the warpPerspective function.
+
+    The parameters threshold_value, min_line_length, and max_line_gap in both methods control the detection of lines in the image. 
+    These parameters affect the number and quality of lines detected, which in turn can impact the accuracy of the perspective correction. 
+    Adjusting these values allows fine-tuning the perspective correction process based on the specific characteristics of the input image. 
+    However, it is important to note that changing these values requires some experimentation to achieve the desired results, and it is recommended to find the optimal values through trial and error.
+    """
+
+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.
+    """

perspective_correction.py

@@ -0,0 +1,106 @@
+import streamlit as st
+import numpy as np
+import cv2
+
+def preprocess_image(image):
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)    
+    threshold = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 15, 7)    
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+    opening = cv2.morphologyEx(threshold, cv2.MORPH_OPEN, kernel, iterations=1)    
+    return opening
+
+def four_point_transform(image, pts):
+    rect = np.zeros((4, 2), dtype=np.float32)
+    s = pts.sum(axis=1)
+    rect[0] = pts[np.argmin(s)]
+    rect[2] = pts[np.argmax(s)]
+    diff = np.diff(pts, axis=1)
+    rect[1] = pts[np.argmin(diff)]
+    rect[3] = pts[np.argmax(diff)]
+    (tl, tr, br, bl) = rect
+    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
+    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
+    maxWidth = max(int(widthA), int(widthB))
+    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
+    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
+    maxHeight = max(int(heightA), int(heightB))
+    dst = np.array([[0, 0], [maxWidth - 1, 0], [maxWidth - 1, maxHeight - 1], [0, maxHeight - 1]], dtype=np.float32)
+    M = cv2.getPerspectiveTransform(rect, dst)
+    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
+
+    return warped
+
+def perspective_correction(image, threshold_value=100, min_line_length=100, max_line_gap=10):
+    processed_image = preprocess_image(image)
+    lines = cv2.HoughLinesP(
+        processed_image,
+        rho=1,
+        theta=np.pi/180,
+        threshold=threshold_value,
+        minLineLength=min_line_length,
+        maxLineGap=max_line_gap
+    )
+
+    if len(lines) < 1:
+        st.write("No lines found.")
+        return image
+
+    endpoints = []
+    for line in lines:
+        x1, y1, x2, y2 = line[0]
+        endpoints.append((x1, y1))
+        endpoints.append((x2, y2))
+
+    if len(endpoints) >= 4:
+        endpoints = np.array(endpoints)
+        corrected_image = four_point_transform(image, endpoints)
+    else:
+        st.write("Insufficient endpoints found.")
+        return image
+
+    return corrected_image
+
+def perspective_correction2(image, threshold_value=100, min_line_length=100, max_line_gap=10):
+    processed_image = preprocess_image(image)
+    lines = cv2.HoughLinesP(
+        processed_image,
+        rho=1,
+        theta=np.pi/180,
+        threshold=threshold_value,
+        minLineLength=min_line_length,
+        maxLineGap=max_line_gap
+    )
+
+    if len(lines) < 1:
+        st.write("No lines found.")
+        return image
+
+    endpoints = []
+    for line in lines:
+        x1, y1, x2, y2 = line[0]
+        endpoints.append((x1, y1))
+        endpoints.append((x2, y2))
+
+    if len(endpoints) >= 4:
+        endpoints = np.array(endpoints)
+        hull = cv2.convexHull(endpoints)
+        rect = np.zeros((4, 2), dtype=np.float32)
+        hull = hull.reshape(-1, 2)
+        rect[0] = hull[np.argmin(hull.sum(axis=1))]
+        rect[2] = hull[np.argmax(hull.sum(axis=1))]
+        rect[1] = hull[np.argmin(np.diff(hull, axis=1))]
+        rect[3] = hull[np.argmax(np.diff(hull, axis=1))]
+        widthA = np.sqrt(((rect[2][0] - rect[3][0]) ** 2) + ((rect[2][1] - rect[3][1]) ** 2))
+        widthB = np.sqrt(((rect[1][0] - rect[0][0]) ** 2) + ((rect[1][1] - rect[0][1]) ** 2))
+        maxWidth = max(int(widthA), int(widthB))
+        heightA = np.sqrt(((rect[1][0] - rect[2][0]) ** 2) + ((rect[1][1] - rect[2][1]) ** 2))
+        heightB = np.sqrt(((rect[0][0] - rect[3][0]) ** 2) + ((rect[0][1] - rect[3][1]) ** 2))
+        maxHeight = max(int(heightA), int(heightB))
+        dst = np.array([[0, 0], [maxWidth - 1, 0], [maxWidth - 1, maxHeight - 1], [0, maxHeight - 1]], dtype=np.float32)
+        M, _ = cv2.findHomography(rect, dst)
+        corrected_image = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
+    else:
+        st.write("Insufficient endpoints found.")
+        return image
+
+    return corrected_image

requirements.txt

@@ -0,0 +1,5 @@
+rembg
+streamlit
+numpy
+streamlit_option_menu
+streamlit-drawable-canvas