Spaces:
Build error
Build error
| # Import libraries | |
| import numpy as np | |
| import streamlit as st | |
| from PIL import ExifTags, Image | |
| from Segmentation.segmentation import get_mask, replace_sofa | |
| from StyleTransfer.styleTransfer import ( | |
| StyleFAST, | |
| StyleTransformer, | |
| styleProjection, | |
| ) | |
| PAGE_CONFIG = { | |
| "page_title": "SofaStyler.io", | |
| "page_icon": ":art:", | |
| "layout": "centered", | |
| } | |
| st.set_page_config(**PAGE_CONFIG) | |
| def fix_orient(img: Image.Image) -> Image.Image: | |
| """ | |
| This function fix the orientation of input images. | |
| This is especially usefull in the context of images from a mobile phone. | |
| Parameters: | |
| img = input image | |
| Return: | |
| img = img with correct orientation | |
| """ | |
| flag = False | |
| for orientation in ExifTags.TAGS.keys(): | |
| if ExifTags.TAGS[orientation] == "Orientation": | |
| flag = True | |
| break | |
| info = img.getexif() | |
| if len(info) & flag: | |
| info = dict(info.items()) | |
| if orientation in info.keys(): | |
| orientation = info[orientation] | |
| if (orientation == 1) | (orientation == 2): | |
| img = img | |
| if (orientation == 3) | (orientation == 4): | |
| img = img.rotate(180, expand=True) | |
| if (orientation == 5) | (orientation == 6): | |
| img = img.rotate(270, expand=True) | |
| if (orientation == 7) | (orientation == 8): | |
| img = img.rotate(90, expand=True) | |
| return img | |
| def resize_sofa(img: Image.Image) -> Image.Image: | |
| """ | |
| This function adds padding to make the original image square and 640by640. | |
| It also returns the original ratio of the image, such that it can be reverted later. | |
| Parameters: | |
| img = original image | |
| Return: | |
| img_square = squared image | |
| box = parameters to later crop the image to it original ratio | |
| """ | |
| width, height = img.size | |
| idx = np.argmin([width, height]) | |
| newsize = (640, 640) # parameters from test script | |
| if idx == 0: | |
| img_square = Image.new(img.mode, (height, height), (255, 255, 255)) | |
| img_square.paste(img, ((height - width) // 2, 0)) | |
| box = ( | |
| newsize[0] * (1 - width / height) // 2, | |
| 0, | |
| newsize[0] - newsize[0] * (1 - width / height) // 2, | |
| newsize[1], | |
| ) | |
| else: | |
| img_square = Image.new(img.mode, (width, width), (255, 255, 255)) | |
| img_square.paste(img, (0, (width - height) // 2)) | |
| box = ( | |
| 0, | |
| newsize[1] * (1 - height / width) // 2, | |
| newsize[0], | |
| newsize[1] - newsize[1] * (1 - height / width) // 2, | |
| ) | |
| img_square = img_square.resize(newsize) | |
| return img_square, box | |
| def resize_style(img: Image.Image) -> Image.Image: | |
| """ | |
| This function generates a zoomed out version of | |
| the style image and resizes it to a 640by640 square. | |
| Parameters: | |
| img = image containing the style/pattern | |
| Return: | |
| dst = a zoomed-out and resized version of the pattern | |
| """ | |
| width, height = img.size | |
| idx = np.argmin([width, height]) | |
| # Makes the image square by cropping | |
| if idx == 0: | |
| top = (height - width) // 2 | |
| bottom = height - (height - width) // 2 | |
| left = 0 | |
| right = width | |
| else: | |
| left = (width - height) // 2 | |
| right = width - (width - height) // 2 | |
| top = 0 | |
| bottom = height | |
| newsize = (640, 640) # parameters from test script | |
| img = img.crop((left, top, right, bottom)) | |
| # Constructs a zoomed-out version | |
| copies = 8 | |
| resize = (newsize[0] // copies, newsize[1] // copies) | |
| img_zoomed_out = Image.new("RGB", (resize[0] * copies, resize[1] * copies)) | |
| img = img.resize((resize)) | |
| for row in range(copies): | |
| img = img.transpose(Image.FLIP_LEFT_RIGHT) | |
| for column in range(copies): | |
| img = img.transpose(Image.FLIP_TOP_BOTTOM) | |
| img_zoomed_out.paste(img, (resize[0] * row, resize[1] * column)) | |
| img_zoomed_out = img_zoomed_out.resize((newsize)) | |
| return img_zoomed_out | |
| image = Image.open("figures/logo.png") # Brand logo image (optional) | |
| options = [ | |
| "Style Transformer", | |
| "StyleFAST", | |
| "Style Projection", | |
| ] | |
| # Create two columns with different width | |
| col1, col2 = st.columns([0.8, 0.2]) | |
| with col1: # To display the header text using css style | |
| st.markdown( | |
| """ | |
| <style> | |
| @import url('https://fonts.googleapis.com/css2?family=Arvo&display=swap'); | |
| </style> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| st.markdown( | |
| """ | |
| <style> | |
| html, body, [class*="css"] { | |
| font-family: 'Arvo'; | |
| } | |
| </style> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| st.markdown( | |
| """ <style> .font { | |
| font-size:30px ; font-family: 'Arvo'; color: #04b188; | |
| src:url("https://fonts.googleapis.com/css2?family=Arvo&display=swap");} | |
| </style> """, | |
| unsafe_allow_html=True, | |
| ) # } </style> """, unsafe_allow_html=True) | |
| st.markdown( | |
| '<p class="font">Upload your photos here...</p>', unsafe_allow_html=True | |
| ) | |
| with col2: # To display brand logo | |
| st.image(image, width=150) | |
| # Add a header and expander in side bar | |
| st.sidebar.markdown('<center class="font">🛋 </center>', unsafe_allow_html=True) | |
| st.sidebar.markdown( | |
| '<center class="font">A sofastyler App</center>', unsafe_allow_html=True | |
| ) | |
| st.sidebar.markdown("") | |
| with st.sidebar.text("About the App"): | |
| st.write( | |
| """ | |
| Customize your sofa to your wildest dreams 💭!\ | |
| \nProvide a picture of your sofa, a desired pattern and\ | |
| choose one of the algorithms below. | |
| \nOr just look at an example. | |
| """ | |
| ) | |
| st.sidebar.title("") | |
| with st.sidebar.expander("References"): | |
| st.write( | |
| "[1. The data that was used to train the segmentation model.]" | |
| + "(https://tianchi.aliyun.com/specials/promotion/alibaba-3d-future)" | |
| + "\n\n" | |
| + "[2. Github repository used to train a segmentation model with transfer " | |
| + "learning.]" | |
| + "(https://github.com/qubvel/segmentation_models)" | |
| + "\n\n" | |
| + "[3. The github repository that is used for the style transformer.]" | |
| + "(https://github.com/diyiiyiii/StyTR-2)" | |
| + "\n\n" | |
| + "[4. A tensorflow model for fast arbitrary image style transfer.]" | |
| + "(https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2)" | |
| + "\n\n" | |
| + "[5. A paddleHub model for parameter free style transfer.]" | |
| + "(https://github.com/PaddlePaddle/PaddleHub/tree/release/v2.2/modules/" | |
| + "image/Image_gan/style_transfer/stylepro_artistic)" | |
| ) | |
| # Add file uploader to allow users to upload photos | |
| uploaded_content = st.file_uploader( | |
| label="Image with sofa", type=["jpg", "png", "jpeg"] | |
| ) | |
| uploaded_style = st.file_uploader( | |
| label="Image with pattern", type=["jpg", "png", "jpeg"] | |
| ) | |
| # Example section | |
| checkbox = st.checkbox("Show example") | |
| if checkbox: | |
| filter = st.radio("Style your sofa with:", options) | |
| col1, col2 = st.columns([0.5, 0.5]) | |
| if filter == "Style Transformer": | |
| content = Image.open("figures/sofa_example1.jpg") | |
| style = Image.open("figures/style_example1.jpg") | |
| output = "figures/0.png" | |
| with col1: | |
| st.markdown( | |
| '<p style="text-align: center;">Before</p>', unsafe_allow_html=True | |
| ) | |
| st.image(content, width=300) | |
| st.image(style, width=300) | |
| with col2: | |
| st.markdown( | |
| '<p style="text-align: center;">After</p>', unsafe_allow_html=True | |
| ) | |
| st.image(output, width=300) | |
| elif filter == "StyleFAST": | |
| content = Image.open("figures/sofa_example3.jpg") | |
| style = Image.open("figures/style_example10.jpg") | |
| output = "figures/1.png" | |
| with col1: | |
| st.markdown( | |
| '<p style="text-align: center;">Before</p>', unsafe_allow_html=True | |
| ) | |
| st.image(content, width=300) | |
| st.image(style, width=300) | |
| with col2: | |
| st.markdown( | |
| '<p style="text-align: center;">After</p>', unsafe_allow_html=True | |
| ) | |
| st.image(output, width=300) | |
| elif filter == "Style Projection": | |
| content = Image.open("figures/sofa_example2.jpg") | |
| style = Image.open("figures/style_example6.jpg") | |
| output = "figures/2.png" | |
| with col1: | |
| st.markdown( | |
| '<p style="text-align: center;">Before</p>', unsafe_allow_html=True | |
| ) | |
| st.image(content, width=300) | |
| st.image(style, width=300) | |
| with col2: | |
| st.markdown( | |
| '<p style="text-align: center;">After</p>', unsafe_allow_html=True | |
| ) | |
| st.image(output, width=300) | |
| # Add 'before' and 'after' columns | |
| elif (uploaded_content is not None) & (uploaded_style is not None): | |
| content = fix_orient(Image.open(uploaded_content)) | |
| style = fix_orient(Image.open(uploaded_style)) | |
| filter = st.radio("Style your sofa with:", options) | |
| ETA = "Unknown" | |
| if filter == "Style Transformer": | |
| ETA = "50s with CPU, 9s with GPU" | |
| elif filter == "StyleFAST": | |
| ETA = "15s with CPU, 3s with GPU" | |
| elif filter == "Style Projection": | |
| alpha = st.slider( | |
| "Adjust the weight of the image vs style", 0.0, 1.0, 0.8, step=0.1 | |
| ) | |
| ETA = "20s with CPU, 10s with GPU" | |
| st.info("Estimated processing time: " + ETA) | |
| button = st.button("Style my sofa") | |
| col1, col2 = st.columns([0.5, 0.5]) | |
| with col1: | |
| st.markdown('<p style="text-align: center;">Before</p>', unsafe_allow_html=True) | |
| st.image(content, width=300) | |
| st.image(style, width=300) | |
| if button: | |
| with col2: | |
| st.markdown( | |
| '<p style="text-align: center;">After</p>', unsafe_allow_html=True | |
| ) | |
| with st.spinner("Preprocessing images..."): | |
| # preprocess input images to be (640,640) squares | |
| # to fit requirements of the segmentation model | |
| resized_img, box = resize_sofa(content) | |
| resized_style = resize_style(style) | |
| # generate mask for image | |
| with st.spinner("generating mask..."): | |
| mask = get_mask(resized_img) | |
| if filter == "Style Transformer": | |
| # Created a styled sofa | |
| with st.spinner("Styling sofa..."): | |
| styled_sofa = StyleTransformer(resized_img, resized_style) | |
| # postprocess the final image | |
| with st.spinner("Replacing sofa..."): | |
| new_sofa = replace_sofa(resized_img, mask, styled_sofa) | |
| new_sofa = new_sofa.crop(box) | |
| st.balloons() | |
| st.image(new_sofa, width=300) | |
| elif filter == "StyleFAST": | |
| # Created a styled sofa | |
| with st.spinner("Styling sofa..."): | |
| styled_sofa = StyleFAST(resized_img, resized_style) | |
| # postprocess the final image | |
| with st.spinner("Replacing sofa..."): | |
| new_sofa = replace_sofa(resized_img, mask, styled_sofa) | |
| new_sofa = new_sofa.crop(box) | |
| st.balloons() | |
| st.image(new_sofa, width=300) | |
| elif filter == "Style Projection": | |
| # Created a styled sofa | |
| with st.spinner("Styling sofa..."): | |
| styled_sofa = styleProjection(resized_img, resized_style, alpha) | |
| # postprocess the final image | |
| with st.spinner("Replacing sofa..."): | |
| new_sofa = replace_sofa(resized_img, mask, styled_sofa) | |
| new_sofa = new_sofa.crop(box) | |
| st.balloons() | |
| st.image(new_sofa, width=300) | |
| else: | |
| st.image(image, width=300) | |
| # Add a feedback section in the sidebar | |
| st.sidebar.title(" ") # create space | |
| st.sidebar.markdown(" ") | |
| st.sidebar.subheader("Please help us improve!") | |
| with st.sidebar.form(key="columns_in_form", clear_on_submit=True): | |
| rating = st.slider( | |
| "Please rate the app", | |
| min_value=1, | |
| max_value=5, | |
| value=3, | |
| help="Drag the slider to rate the app." | |
| + "This is a 1-5 rating scale where 5 is the highest rating", | |
| ) | |
| text = st.text_input(label="Please leave your feedback here") | |
| submitted = st.form_submit_button("Submit") | |
| if submitted: | |
| st.write("Thanks for your feedback!") | |
| st.markdown("Your Rating:") | |
| st.markdown(rating) | |
| st.markdown("Your Feedback:") | |
| st.markdown(text) | |