added Save image code

Pages/bgImages.py

@@ -43,7 +43,7 @@ def BGIMAGES():
         st.image(image, caption=f"Generated image for: {prompt}", width = 300)
 
     else:
-        st.write("Please enter a movie title to generate an image.")
+        st.write("Please enter prompt for background to generate an image.")
 
     # Path to the folder containing images
     image_folder = "Assets/bgImages"

Pages/createdVisuals.py

@@ -8,6 +8,8 @@ def load_images_from_folder(folder):
             images.append(os.path.join(folder, filename))
     return images
 
+def get_image_name(image_path):
+    return os.path.splitext(os.path.basename(image_path))[0]
 
 # Main function
 def CREATEDIMAGES():
@@ -20,12 +22,12 @@ def CREATEDIMAGES():
     images = load_images_from_folder(image_folder)
 
     # Display images and information in a grid layout with three images per row
-    col_width = 350 # Adjust this value according to your preference
+    col_width = 350  # Adjust this value according to your preference
     num_images = len(images)
     images_per_row = 3
     num_rows = (num_images + images_per_row - 1) // images_per_row
 
-    st.header('Available Templates', divider = 'red')
+    st.header('Available Templates', divider='red')
 
     # Display images and information in a grid layout with three images per row
     for i in range(num_rows):
@@ -33,11 +35,9 @@ def CREATEDIMAGES():
         for j in range(images_per_row):
             idx = i * images_per_row + j
             if idx < num_images:
+                image_name = get_image_name(images[idx])
                 cols[j].image(images[idx], width=col_width)
-                cols[j].write(images[idx])
-
+                cols[j].write(image_name)
 
 if __name__ == "__main__":
     CREATEDIMAGES()
-
-

Pages/imageBB.py

@@ -9,8 +9,7 @@ import cv2
 import numpy as np
 
 @st.cache_data
-def get_masks(rect,img_path):
-
+def get_masks(rect, img_path):
     CHECKPOINT_PATH = os.path.join("weights", "sam_vit_h_4b8939.pth")
     DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
     MODEL_TYPE = "vit_h"
@@ -34,8 +33,6 @@ def get_masks(rect,img_path):
     )
     return masks
 
-
-
 def run(img_path):
     st.set_option("deprecation.showfileUploaderEncoding", False)
 
@@ -51,14 +48,11 @@ def run(img_path):
         rects = st_img_label(resized_img, box_color="red", rects=st.session_state.rects)
         st.session_state.rects = rects
     else:
-        st.image(resized_img, caption="Uploaded Image",width=300, use_column_width=True)
+        st.image(resized_img, caption="Uploaded Image", width=300, use_column_width=True)
 
         for rect in st.session_state.rects:
-            # st.write(f"Rectangle: {rect}")
-
             with st.spinner('Please wait while the product image is being extracted...'):
-
-                masks = get_masks(rect,img_path)
+                masks = get_masks(rect, img_path)
 
                 save_dir = "Assets/saved_masks"
                 if not os.path.exists(save_dir):
@@ -84,81 +78,85 @@ def run(img_path):
                 st.image(image, width=300, caption=f"Selected image: {selected_image}")
 
                 if st.button("Create Image"):
-                    # Read the base image and background image
-                    image_bgr = cv2.imread(img_path)
-                    background_bgr = cv2.imread(image_pathBG)
-                    
-                    # Resize the background image to match the size of image_bgr
-                    background_bgr = cv2.resize(background_bgr, (image_bgr.shape[1], image_bgr.shape[0]))
-
-                    # Convert the base image to RGB format for mask prediction if it's not already in RGB
-                    if image_bgr.shape[2] == 3:  # No alpha channel, standard BGR image
-                        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
-                    else:
-                        image_rgb = image_bgr[:, :, :3]  # Drop alpha channel if it exists
-
-                    # Assuming masks is a binary mask, convert it to uint8 format
-                    mask = (masks[0] > 0).astype(np.uint8) * 255
-
-                    # Apply a Gaussian blur to the mask to smooth the edges
-                    mask = cv2.GaussianBlur(mask, (3,3), 0)
-
-                    # Ensure the image has an alpha channel
-                    if image_bgr.shape[2] == 3:  # If no alpha channel, add one
-                        b, g, r = cv2.split(image_bgr)
-                        alpha_channel = mask  # Use the blurred mask as the alpha channel
-                        image_bgra = cv2.merge((b, g, r, alpha_channel))
-                    else:
-                        image_bgra = image_bgr
-
-                    # Get the dimensions of the images
-                    masked_height, masked_width = image_bgra.shape[:2]
-                    background_height, background_width = background_bgr.shape[:2]
-
-                    # Calculate the coordinates to place the masked image in the center of the background image
-                    x_offset = (background_width - masked_width) // 2
-                    y_offset = (background_height - masked_height) // 2
-
-                    # Resize the masked image if it is larger than the background area
-                    if masked_width > background_width or masked_height > background_height:
-                        scaling_factor = min(background_width / masked_width, background_height / masked_height)
-                        new_size = (int(masked_width * scaling_factor), int(masked_height * scaling_factor))
-                        image_bgra = cv2.resize(image_bgra, new_size, interpolation=cv2.INTER_AREA)
-                        masked_height, masked_width = image_bgra.shape[:2]
-                        x_offset = (background_width - masked_width) // 2
-                        y_offset = (background_height - masked_height) // 2
-
-                    # Create a copy of the background image and convert it to BGRA
-                    background_bgra = cv2.cvtColor(background_bgr, cv2.COLOR_BGR2BGRA)
-
-                    # Overlay the masked image onto the center of the background image
-                    overlay_image = background_bgra.copy()
-
-                    # Only update the region where the segmented image will be placed
-                    overlay = np.zeros_like(background_bgra)
-                    overlay[y_offset:y_offset+masked_height, x_offset:x_offset+masked_width] = image_bgra
-
-                    # Create the alpha mask for blending
-                    alpha_mask = overlay[:, :, 3] / 255.0
-                    alpha_inv = 1.0 - alpha_mask
-
-                    # Modify alpha channel for smoother blending
-                    alpha_mask = alpha_mask ** 0.5  # Applying square root for smoother blending
-
-                    # Blend the images
-                    for c in range(0, 3):
-                        overlay_image[:, :, c] = (alpha_mask * overlay[:, :, c] + alpha_inv * overlay_image[:, :, c])
-
-                    # Set the alpha channel
-                    overlay_image[:, :, 3] = np.clip(overlay[:, :, 3] + background_bgra[:, :, 3], 0, 255)
-
-                    # Save the result
-                    output_path = 'Assets/output/images/overlay_image.png'
-                    cv2.imwrite(output_path, overlay_image)
-
-                    # Display the overlay image
-                    st.image(output_path, caption="Overlay Image", use_column_width=True, width=300)
-
+                    st.session_state.create_image = True
+
+        if st.session_state.get("create_image"):
+            # Read the base image and background image
+            image_bgr = cv2.imread(img_path)
+            background_bgr = cv2.imread(image_pathBG)
+
+            # Resize the background image to match the size of image_bgr
+            background_bgr = cv2.resize(background_bgr, (image_bgr.shape[1], image_bgr.shape[0]))
+
+            # Convert the base image to RGB format for mask prediction if it's not already in RGB
+            if image_bgr.shape[2] == 3:  # No alpha channel, standard BGR image
+                image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+            else:
+                image_rgb = image_bgr[:, :, :3]  # Drop alpha channel if it exists
+
+            # Assuming masks is a binary mask, convert it to uint8 format
+            mask = (masks[0] > 0).astype(np.uint8) * 255
+
+            # Apply a Gaussian blur to the mask to smooth the edges
+            mask = cv2.GaussianBlur(mask, (3, 3), 0)
+
+            # Ensure the image has an alpha channel
+            if image_bgr.shape[2] == 3:  # If no alpha channel, add one
+                b, g, r = cv2.split(image_bgr)
+                alpha_channel = mask  # Use the blurred mask as the alpha channel
+                image_bgra = cv2.merge((b, g, r, alpha_channel))
+            else:
+                image_bgra = image_bgr
+
+            # Get the dimensions of the images
+            masked_height, masked_width = image_bgra.shape[:2]
+            background_height, background_width = background_bgr.shape[:2]
+
+            # Calculate the coordinates to place the masked image in the center of the background image
+            x_offset = (background_width - masked_width) // 2
+            y_offset = (background_height - masked_height) // 2
+
+            # Resize the masked image if it is larger than the background area
+            if masked_width > background_width or masked_height > background_height:
+                scaling_factor = min(background_width / masked_width, background_height / masked_height)
+                new_size = (int(masked_width * scaling_factor), int(masked_height * scaling_factor))
+                image_bgra = cv2.resize(image_bgra, new_size, interpolation=cv2.INTER_AREA)
+                masked_height, masked_width = image_bgra.shape[:2]
+                x_offset = (background_width - masked_width) // 2
+                y_offset = (background_height - masked_height) // 2
+
+            # Create a copy of the background image and convert it to BGRA
+            background_bgra = cv2.cvtColor(background_bgr, cv2.COLOR_BGR2BGRA)
+
+            # Overlay the masked image onto the center of the background image
+            overlay_image = background_bgra.copy()
+
+            # Only update the region where the segmented image will be placed
+            overlay = np.zeros_like(background_bgra)
+            overlay[y_offset:y_offset + masked_height, x_offset:x_offset + masked_width] = image_bgra
+
+            # Create the alpha mask for blending
+            alpha_mask = overlay[:, :, 3] / 255.0
+            alpha_inv = 1.0 - alpha_mask
+
+            # Modify alpha channel for smoother blending
+            alpha_mask = alpha_mask ** 0.5  # Applying square root for smoother blending
+
+            # Blend the images
+            for c in range(0, 3):
+                overlay_image[:, :, c] = (alpha_mask * overlay[:, :, c] + alpha_inv * overlay_image[:, :, c])
+
+            # Set the alpha channel
+            overlay_image[:, :, 3] = np.clip(overlay[:, :, 3] + background_bgra[:, :, 3], 0, 255)
+
+            # Prompt user for the filename
+            filename = st.text_input("Enter a name to save the image:")
+            if filename and st.button("Save Image"):
+                output_path = f'Assets/output/images/{filename}.png'
+                cv2.imwrite(output_path, overlay_image)
+
+                # Display the overlay image
+                st.image(output_path, caption="Created Image", use_column_width=True, width=300)
 
     def annotate():
         st.session_state.saved = True
@@ -166,3 +164,5 @@ def run(img_path):
     if st.session_state.rects:
         st.button(label="Save", on_click=annotate)
 
+# Example of calling the function
+# run("path/to/your/image.jpg")