ir12345's picture
Upload app.py
a1b7adb verified
import cv2
import numpy as np
import io
import streamlit as st
from PIL import Image
def color_transfer(source, target):
# Convert images to LAB color space
source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype("float32")
target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype("float32")
# Compute mean and std for each channel
mean_src, std_src = cv2.meanStdDev(source)
mean_tgt, std_tgt = cv2.meanStdDev(target)
# Reshape for broadcasting
mean_src, std_src = mean_src.reshape((1, 1, 3)), std_src.reshape((1, 1, 3))
mean_tgt, std_tgt = mean_tgt.reshape((1, 1, 3)), std_tgt.reshape((1, 1, 3))
# Apply color transfer
adjusted = ((target - mean_tgt) * (std_src / std_tgt)) + mean_src
adjusted = np.clip(adjusted, 0, 255).astype("uint8")
return cv2.cvtColor(adjusted, cv2.COLOR_LAB2BGR)
def main():
st.set_page_config(layout="wide")
st.sidebar.header("Upload Images")
# Upload two images in the sidebar
uploaded_file1 = st.sidebar.file_uploader("Upload first image", type=["jpg", "png", "jpeg"])
uploaded_file2 = st.sidebar.file_uploader("Upload second image", type=["jpg", "png", "jpeg"])
st.sidebar.image("test.png", width=700)
# Display the images in the center
col1, col2, col3 = st.columns([1, 1, 1])
with col1:
st.header("Uploaded Images")
if uploaded_file1 and uploaded_file2:
ref_img = Image.open(uploaded_file1)
input_img = Image.open(uploaded_file2)
# Display images side by side
st.image([ref_img, input_img], caption=["Reference Image", "Input Image"], width=300)
elif uploaded_file1:
ref_img = Image.open(uploaded_file1)
st.image(ref_img, caption="Reference Image", width=300)
elif uploaded_file2:
input_img = Image.open(uploaded_file2)
st.image(input_img, caption="Input Image", width=300)
else:
st.write("Upload images to display them here.")
with col2:
st.header(" ")
if uploaded_file1 and uploaded_file2:
img1 = Image.open(uploaded_file1)
img2 = Image.open(uploaded_file2)
ref_img = cv2.cvtColor(np.array(img1), cv2.COLOR_RGB2BGR)
input_img = cv2.cvtColor(np.array(img2), cv2.COLOR_RGB2BGR)
orig_h, orig_w = input_img.shape[:2]
input_img = cv2.resize(input_img, (ref_img.shape[1], ref_img.shape[0]))
# Apply color transfer
output_img = color_transfer(ref_img, input_img)
output_img_resize = cv2.resize(output_img, (orig_w, orig_h))
st.image(output_img_resize, caption="Result", width=500, channels="BGR")
_, buffer = cv2.imencode(".jpg", output_img_resize)
img_bytes = io.BytesIO(buffer)
# Add Download Button
st.download_button(
label="Download Image",
data=img_bytes,
file_name="image.jpg",
mime="image/jpeg"
)
if __name__ == "__main__":
main()