FLODARELTIH commited on
Commit
fddb466
·
verified ·
1 Parent(s): f0f6622

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +239 -0
app.py ADDED
@@ -0,0 +1,239 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app.py
2
+ import cv2
3
+ import insightface
4
+ import numpy as np
5
+ from flask import Flask, request, jsonify
6
+ from flask_cors import CORS
7
+ import base64
8
+ import os
9
+ import tempfile
10
+ import uuid
11
+ from werkzeug.utils import secure_filename
12
+ from insightface.app import FaceAnalysis
13
+
14
+ app = Flask(__name__)
15
+ CORS(app)
16
+
17
+ # Configuration
18
+ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB max file size
19
+ app.config['UPLOAD_FOLDER'] = tempfile.gettempdir()
20
+ ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp'}
21
+
22
+ # Initialize face analysis globally
23
+ face_app = FaceAnalysis(providers=['CPUExecutionProvider'])
24
+ face_app.prepare(ctx_id=0, det_size=(640, 640))
25
+
26
+ # Initialize face swapper
27
+ swapper = insightface.model_zoo.get_model('inswapper_128.onnx')
28
+
29
+ def allowed_file(filename):
30
+ return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
31
+
32
+ def decode_base64_image(base64_string):
33
+ """Decode base64 string to numpy image"""
34
+ if ',' in base64_string:
35
+ base64_string = base64_string.split(',')[1]
36
+ img_data = base64.b64decode(base64_string)
37
+ nparr = np.frombuffer(img_data, np.uint8)
38
+ return cv2.imdecode(nparr, cv2.IMREAD_COLOR)
39
+
40
+ def encode_image_to_base64(image):
41
+ """Encode numpy image to base64 string"""
42
+ _, buffer = cv2.imencode('.jpg', image)
43
+ return base64.b64encode(buffer).decode('utf-8')
44
+
45
+ @app.route('/health', methods=['GET'])
46
+ def health_check():
47
+ return jsonify({'status': 'healthy', 'message': 'Face Swap API is running'})
48
+
49
+ @app.route('/swap', methods=['POST'])
50
+ def swap_faces():
51
+ """
52
+ Swap faces between source and target images
53
+ Expected JSON payload:
54
+ {
55
+ "source_image": "base64_string or file_path",
56
+ "target_image": "base64_string or file_path",
57
+ "source_is_base64": true, # optional, defaults to false
58
+ "target_is_base64": true # optional, defaults to false
59
+ }
60
+ Or use multipart/form-data with files:
61
+ - source_file (image file)
62
+ - target_file (image file)
63
+ """
64
+
65
+ try:
66
+ source_img = None
67
+ target_img = None
68
+
69
+ # Handle JSON request
70
+ if request.is_json:
71
+ data = request.get_json()
72
+
73
+ # Get source image
74
+ source_is_base64 = data.get('source_is_base64', False)
75
+ if source_is_base64:
76
+ source_img = decode_base64_image(data['source_image'])
77
+ else:
78
+ source_img = cv2.imread(data['source_image'])
79
+
80
+ # Get target image
81
+ target_is_base64 = data.get('target_is_base64', False)
82
+ if target_is_base64:
83
+ target_img = decode_base64_image(data['target_image'])
84
+ else:
85
+ target_img = cv2.imread(data['target_image'])
86
+
87
+ # Handle multipart form request
88
+ elif 'source_file' in request.files and 'target_file' in request.files:
89
+ source_file = request.files['source_file']
90
+ target_file = request.files['target_file']
91
+
92
+ if source_file and allowed_file(source_file.filename):
93
+ source_filename = secure_filename(source_file.filename)
94
+ source_path = os.path.join(app.config['UPLOAD_FOLDER'], f"source_{uuid.uuid4()}_{source_filename}")
95
+ source_file.save(source_path)
96
+ source_img = cv2.imread(source_path)
97
+ os.remove(source_path) # Clean up
98
+ else:
99
+ return jsonify({'error': 'Invalid source file type'}), 400
100
+
101
+ if target_file and allowed_file(target_file.filename):
102
+ target_filename = secure_filename(target_file.filename)
103
+ target_path = os.path.join(app.config['UPLOAD_FOLDER'], f"target_{uuid.uuid4()}_{target_filename}")
104
+ target_file.save(target_path)
105
+ target_img = cv2.imread(target_path)
106
+ os.remove(target_path) # Clean up
107
+ else:
108
+ return jsonify({'error': 'Invalid target file type'}), 400
109
+
110
+ else:
111
+ return jsonify({'error': 'Invalid request. Provide source_image/target_image or source_file/target_file'}), 400
112
+
113
+ # Validate images
114
+ if source_img is None:
115
+ return jsonify({'error': 'Could not read source image'}), 400
116
+ if target_img is None:
117
+ return jsonify({'error': 'Could not read target image'}), 400
118
+
119
+ # Detect faces
120
+ source_faces = face_app.get(source_img)
121
+ target_faces = face_app.get(target_img)
122
+
123
+ if len(source_faces) == 0:
124
+ return jsonify({'error': 'No face found in source image'}), 400
125
+ if len(target_faces) == 0:
126
+ return jsonify({'error': 'No face found in target image'}), 400
127
+
128
+ # Perform face swap
129
+ swapped_image = swapper.get(target_img, target_faces[0], source_faces[0], paste_back=True)
130
+
131
+ # Prepare response
132
+ return_type = request.args.get('return_type', 'base64')
133
+
134
+ if return_type == 'file':
135
+ # Save and return file path
136
+ output_filename = f"swapped_{uuid.uuid4()}.jpg"
137
+ output_path = os.path.join(app.config['UPLOAD_FOLDER'], output_filename)
138
+ cv2.imwrite(output_path, swapped_image)
139
+
140
+ return jsonify({
141
+ 'success': True,
142
+ 'message': 'Face swap completed successfully',
143
+ 'output_path': output_path,
144
+ 'filename': output_filename
145
+ })
146
+ else:
147
+ # Return base64 encoded image
148
+ encoded_image = encode_image_to_base64(swapped_image)
149
+ return jsonify({
150
+ 'success': True,
151
+ 'message': 'Face swap completed successfully',
152
+ 'swapped_image': encoded_image
153
+ })
154
+
155
+ except Exception as e:
156
+ return jsonify({'error': str(e)}), 500
157
+
158
+ @app.route('/swap/batch', methods=['POST'])
159
+ def batch_swap_faces():
160
+ """
161
+ Batch face swap with multiple targets
162
+ Expected JSON payload:
163
+ {
164
+ "source_image": "base64_string or path",
165
+ "target_images": ["base64_string1", "base64_string2", ...],
166
+ "source_is_base64": true,
167
+ "targets_are_base64": true
168
+ }
169
+ """
170
+ try:
171
+ data = request.get_json()
172
+
173
+ # Get source image
174
+ source_is_base64 = data.get('source_is_base64', False)
175
+ if source_is_base64:
176
+ source_img = decode_base64_image(data['source_image'])
177
+ else:
178
+ source_img = cv2.imread(data['source_image'])
179
+
180
+ if source_img is None:
181
+ return jsonify({'error': 'Could not read source image'}), 400
182
+
183
+ # Detect source face once
184
+ source_faces = face_app.get(source_img)
185
+ if len(source_faces) == 0:
186
+ return jsonify({'error': 'No face found in source image'}), 400
187
+
188
+ source_face = source_faces[0]
189
+
190
+ # Process all target images
191
+ results = []
192
+ target_images = data.get('target_images', [])
193
+ targets_are_base64 = data.get('targets_are_base64', False)
194
+
195
+ for idx, target_img_data in enumerate(target_images):
196
+ try:
197
+ if targets_are_base64:
198
+ target_img = decode_base64_image(target_img_data)
199
+ else:
200
+ target_img = cv2.imread(target_img_data)
201
+
202
+ if target_img is None:
203
+ results.append({'index': idx, 'error': 'Could not read target image'})
204
+ continue
205
+
206
+ target_faces = face_app.get(target_img)
207
+ if len(target_faces) == 0:
208
+ results.append({'index': idx, 'error': 'No face found in target image'})
209
+ continue
210
+
211
+ swapped_image = swapper.get(target_img, target_faces[0], source_face, paste_back=True)
212
+ encoded_image = encode_image_to_base64(swapped_image)
213
+
214
+ results.append({
215
+ 'index': idx,
216
+ 'success': True,
217
+ 'swapped_image': encoded_image
218
+ })
219
+ except Exception as e:
220
+ results.append({'index': idx, 'error': str(e)})
221
+
222
+ return jsonify({
223
+ 'success': True,
224
+ 'message': f'Processed {len(results)} images',
225
+ 'results': results
226
+ })
227
+
228
+ except Exception as e:
229
+ return jsonify({'error': str(e)}), 500
230
+
231
+ if __name__ == '__main__':
232
+ print("Starting Face Swap API Server...")
233
+ print("Make sure 'inswapper_128.onnx' is in the current directory")
234
+ print("API endpoints:")
235
+ print(" GET /health - Health check")
236
+ print(" POST /swap - Single face swap")
237
+ print(" POST /swap/batch - Batch face swap")
238
+ print("\nStarting server on http://localhost:5000")
239
+ app.run(host='0.0.0.0', port=7860, debug=True)