Spaces:
Runtime error
Runtime error
File size: 2,691 Bytes
bbb8961 |
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 70 71 72 73 |
import streamlit as st
import face_recognition
import os
import cv2
import insightface
import pickle
from insightface.app import FaceAnalysis
from concurrent.futures import ThreadPoolExecutor
# Initialize your models only once
app = FaceAnalysis(name='buffalo_l')
app.prepare(ctx_id=0, det_size=(640, 640))
swapper = insightface.model_zoo.get_model(os.path.join(os.getcwd(), 'inswapper_128.onnx'), download=False)
# Load pickle file once and keep it in memory
@st.cache(allow_output_mutation=True)
def load_data_images():
with open(os.path.join(os.getcwd(), 'data_images.pkl'), 'rb') as file:
return pickle.load(file)
data_images = load_data_images()
def face_swapper(image_background, image_customer):
face_customer = app.get(image_customer)[0]
faces = app.get(image_background)
for face in faces:
image_background = swapper.get(image_background, face, face_customer, paste_back=True)
return image_background
def process(image):
images_background_encoding, images_background_contents = data_images['encodings'], data_images['content']
image_loaded = face_recognition.load_image_file(image)
face_encoding = face_recognition.face_encodings(image_loaded)[0]
face_distances = face_recognition.face_distance(images_background_encoding, face_encoding)
tmp_distance = face_distances[0]
tmp_content = images_background_contents[0]
for face_distance, images_background_content in zip(face_distances[1:], images_background_contents[1:]):
if tmp_distance > face_distance:
tmp_distance = face_distance
tmp_content = images_background_content
output_image = face_swapper(tmp_content, image_loaded)
return output_image
image_output = None
st.title('Change Faces')
option = st.radio('How would you like to upload your image?', ('File', 'WebCam'), horizontal=True)
if option=='File':
uploaded_file = st.file_uploader('Choose your image', type=['jpg', 'png', 'jpeg'])
else:
uploaded_file = st.camera_input("Take a picture")
if uploaded_file is not None:
bytes_data = uploaded_file.getvalue()
if option=='File':
st.image(uploaded_file)
if st.button('Process'):
with ThreadPoolExecutor() as executor:
future = executor.submit(process, uploaded_file)
image_output = future.result()
st.image(image_output)
if image_output is not None:
image_output_to_download = cv2.cvtColor(image_output, cv2.COLOR_BGR2RGB)
_, image_output_to_download = cv2.imencode('.jpg', image_output_to_download)
st.download_button('Download image', image_output_to_download.tobytes(), file_name=f'output_{uploaded_file.name}') |