app2.py

import streamlit as st
import face_recognition
import os
import cv2
import insightface
import pickle

from insightface.app import FaceAnalysis


PATH_TMP = 'tmp'
PATH_MODEL = 'bin'
DATA_IMAGE_PICKLE = 'data_images.pkl'
ONNX_SWAPPER_MODEL = 'inswapper_128.onnx'


def face_swapper(image_background, image_customer):
 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(), ONNX_SWAPPER_MODEL), download=False)
 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):
 with open(os.path.join(os.getcwd(), DATA_IMAGE_PICKLE), 'rb') as file:
 data_images = pickle.load(file)

 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'):
 image_output = process(uploaded_file)
 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}')