sandrocalzada commited on
Commit
d29a232
1 Parent(s): fedaa13

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +71 -0
app.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import face_recognition
3
+ import os
4
+ import cv2
5
+ import insightface
6
+ import pickle
7
+
8
+ from insightface.app import FaceAnalysis
9
+
10
+
11
+ PATH_TMP = 'tmp'
12
+ PATH_MODEL = 'bin'
13
+ DATA_IMAGE_PICKLE = 'data_images.pkl'
14
+ ONNX_SWAPPER_MODEL = 'inswapper_128.onnx'
15
+
16
+
17
+ def face_swapper(image_background, image_customer):
18
+ app = FaceAnalysis(name='buffalo_l')
19
+ app.prepare(ctx_id=0, det_size=(640, 640))
20
+ swapper = insightface.model_zoo.get_model(os.path.join(os.getcwd(), PATH_MODEL, ONNX_SWAPPER_MODEL), download=False)
21
+ face_customer = app.get(image_customer)[0]
22
+ faces = app.get(image_background)
23
+
24
+ for face in faces:
25
+ image_background = swapper.get(image_background, face, face_customer, paste_back=True)
26
+
27
+ return image_background
28
+
29
+ def process(image):
30
+ with open(os.path.join(os.getcwd(), PATH_MODEL, DATA_IMAGE_PICKLE), 'rb') as file:
31
+ data_images = pickle.load(file)
32
+
33
+ images_background_encoding, images_background_contents = data_images['encodings'], data_images['content']
34
+ image_loaded = face_recognition.load_image_file(image)
35
+ face_encoding = face_recognition.face_encodings(image_loaded)[0]
36
+ face_distances = face_recognition.face_distance(images_background_encoding, face_encoding)
37
+
38
+ tmp_distance = face_distances[0]
39
+ tmp_content = images_background_contents[0]
40
+ for face_distance, images_background_content in zip(face_distances[1:], images_background_contents[1:]):
41
+ if tmp_distance > face_distance:
42
+ tmp_distance = face_distance
43
+ tmp_content = images_background_content
44
+
45
+ output_image = face_swapper(tmp_content, image_loaded)
46
+ return output_image
47
+
48
+ image_output = None
49
+
50
+ st.title('Change Faces')
51
+
52
+ option = st.radio('How would you like to upload your image?', ('File', 'WebCam'), horizontal=True)
53
+
54
+ if option=='File':
55
+ uploaded_file = st.file_uploader('Choose your image', type=['jpg', 'png', 'jpeg'])
56
+ else:
57
+ uploaded_file = st.camera_input("Take a picture")
58
+
59
+
60
+ if uploaded_file is not None:
61
+ bytes_data = uploaded_file.getvalue()
62
+ if option=='File':
63
+ st.image(uploaded_file)
64
+ if st.button('Process'):
65
+ image_output = process(uploaded_file)
66
+ st.image(image_output)
67
+
68
+ if image_output is not None:
69
+ image_output_to_download = cv2.cvtColor(image_output, cv2.COLOR_BGR2RGB)
70
+ _, image_output_to_download = cv2.imencode('.jpg', image_output_to_download)
71
+ st.download_button('Download image', image_output_to_download.tobytes(), file_name=f'output_{uploaded_file.name}')