Spaces:
Sleeping
Sleeping
File size: 2,638 Bytes
c0e7e20 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import os
import cv2
import streamlit as st
import numpy as np
from face_detection import select_face, select_all_faces
from face_swap import face_swap
# Function to read image and convert to RGB
def read_image(image_file):
image = np.array(bytearray(image_file.read()), dtype=np.uint8)
image = cv2.imdecode(image, cv2.IMREAD_COLOR)
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Streamlit app
st.title('FaceSwapApp')
st.sidebar.title("Upload Images")
src_file = st.sidebar.file_uploader("Upload Source Image", type=["jpg", "jpeg", "png"])
dst_file = st.sidebar.file_uploader("Upload Target Image", type=["jpg", "jpeg", "png"])
warp_2d = st.sidebar.checkbox('2D Warp', value=False)
correct_color = st.sidebar.checkbox('Correct Color', value=False)
submit = st.sidebar.button('Generate')
if submit:
if src_file and dst_file:
# Read and display the uploaded images
src_image = read_image(src_file)
dst_image = read_image(dst_file)
# Set up the layout
col1, col2 = st.columns([1, 1])
with col1:
st.image(src_image, caption='Source Image', use_column_width=True)
st.image(dst_image, caption='Target Image', use_column_width=True)
# Process the images
src_points, src_shape, src_face = select_face(src_image)
dst_faceBoxes = select_all_faces(dst_image)
if dst_faceBoxes is None:
st.error('No face detected in target image!')
else:
output = dst_image.copy()
for k, dst_face in dst_faceBoxes.items():
output = face_swap(src_face, dst_face["face"], src_points,
dst_face["points"], dst_face["shape"],
output, warp_2d=warp_2d, correct_color=correct_color)
# Display the result
with col2:
st.image(output, caption='Swapped Image', use_column_width=True)
# Option to download the result
result_path = 'output/result.png'
cv2.imwrite(result_path, cv2.cvtColor(output, cv2.COLOR_RGB2BGR))
with open(result_path, "rb") as file:
st.download_button(
label="Download Swapped Image",
data=file,
file_name="result.png",
mime="image/png"
)
os.remove(result_path)
else:
st.error('Please upload both source and target images!')
|