app.py

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 to None
app = None
swapper = None
data_images = None

# Load pickle file once and keep it in memory
@st.cache_data
def load_data_images():
 with open(os.path.join(os.getcwd(), 'data_images.pkl'), 'rb') as file:
 return pickle.load(file)

# Lazy initialization of data_images
def get_data_images():
 global data_images
 if data_images is None:
 data_images = load_data_images()
 return data_images

# Lazy initialization for models
def prepare_models():
 global app, swapper
 if app is None:
 app = FaceAnalysis(name='buffalo_l')
 app.prepare(ctx_id=0, det_size=(640, 640))
 if swapper is None:
 swapper = insightface.model_zoo.get_model(os.path.join(os.getcwd(), 'inswapper_128.onnx'), download=False)

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):
 prepare_models() # Prepare models only when required
 images_background_encoding, images_background_contents = get_data_images()['encodings'], get_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 st.spinner('Processing...'):
 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}')