Spaces:
Running
Running
from segment_anything import sam_model_registry, SamPredictor | |
import torch | |
import streamlit as st | |
from Pages.streamlit_img_label import st_img_label | |
from Pages.streamlit_img_label.manage import ImageManager | |
import os | |
from PIL import Image | |
import cv2 | |
import numpy as np | |
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" | |
sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE) | |
mask_predictor = SamPredictor(sam) | |
rect = np.array([ | |
rect['left'], | |
rect['top'], | |
rect['left'] + rect['width'], | |
rect['top'] + rect['height'] | |
]) | |
image_bgr = cv2.imread(img_path) | |
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) | |
mask_predictor.set_image(image_rgb) | |
masks, scores, logits = mask_predictor.predict( | |
box=rect, | |
multimask_output=False | |
) | |
return masks | |
def run(img_path): | |
st.set_option("deprecation.showfileUploaderEncoding", False) | |
im = ImageManager(img_path) | |
resized_img = im.resizing_img() | |
resized_rects = im.get_resized_rects() | |
if "rects" not in st.session_state: | |
st.session_state.rects = resized_rects | |
# Only display st_img_label if Save button hasn't been clicked | |
if not st.session_state.get("saved"): | |
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) | |
for rect in st.session_state.rects: | |
with st.spinner('Please wait while the product image is being extracted...'): | |
masks = get_masks(rect, img_path) | |
save_dir = "Assets/saved_masks" | |
if not os.path.exists(save_dir): | |
os.makedirs(save_dir) | |
for i, mask in enumerate(masks): | |
inverted_mask = 255 - (mask * 255).astype(np.uint8) | |
file_path = os.path.join(save_dir, f"inverted_mask_{i}.png") | |
cv2.imwrite(file_path, inverted_mask) | |
print(f"Inverted masks saved to directory: {save_dir}") | |
image_files = [f for f in os.listdir("Assets/bgImages") if os.path.isfile(os.path.join("Assets/bgImages", f))] | |
st.header("Template Selection") | |
# Create a dropdown with the list of image files | |
selected_image = st.selectbox("Select an image file", image_files) | |
if selected_image: | |
# Display the selected image | |
image_pathBG = os.path.join("Assets/bgImages", selected_image) | |
image = Image.open(image_pathBG) | |
st.image(image, width=300, caption=f"Selected image: {selected_image}") | |
if st.button("Create Image"): | |
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 | |
if st.session_state.rects: | |
st.button(label="Save", on_click=annotate) | |
# Example of calling the function | |
# run("path/to/your/image.jpg") | |