import os import cv2 import shutil import time import base64 import json # For parsing target_ids_json from flask import Flask, render_template, request, jsonify, send_from_directory, Response from flask_cors import CORS from werkzeug.utils import secure_filename import modules.globals as Globals import modules.core as core from modules.utilities import normalize_output_path, get_temp_directory_path, is_image as util_is_image from modules.face_analyser import get_one_face, get_many_faces, get_unique_faces_from_target_image, simplify_maps # Added simplify_maps import modules.processors.frame.core as frame_processors_core VIDEO_CAMERA = None target_path_web = None prev_time = 0 frame_count = 0 # For FPS calculation current_fps = 0 # For FPS calculation # Attempt to load initial settings from a file if it exists # This is a placeholder for more sophisticated settings management. # For now, we rely on defaults in modules.globals or explicit setting via UI. # if os.path.exists('switch_states.json'): # try: # with open('switch_states.json', 'r') as f: # import json # states = json.load(f) # # Assuming states directly map to Globals attributes # for key, value in states.items(): # if hasattr(Globals, key): # setattr(Globals, key, value) # except Exception as e: # print(f"Error loading switch_states.json: {e}") app = Flask(__name__) CORS(app) # Enable CORS for all routes UPLOAD_FOLDER = os.path.join(os.getcwd(), 'uploads') PROCESSED_OUTPUTS_FOLDER = os.path.join(os.getcwd(), 'processed_outputs') # Added if not os.path.exists(UPLOAD_FOLDER): os.makedirs(UPLOAD_FOLDER) if not os.path.exists(PROCESSED_OUTPUTS_FOLDER): # Added os.makedirs(PROCESSED_OUTPUTS_FOLDER) app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER app.config['PROCESSED_OUTPUTS_FOLDER'] = PROCESSED_OUTPUTS_FOLDER # Added @app.route('/') def index(): # Renamed from hello_world return render_template('index.html') @app.route('/upload/source', methods=['POST']) def upload_source(): if 'file' not in request.files: return jsonify({'error': 'No file part'}), 400 file = request.files['file'] if file.filename == '': return jsonify({'error': 'No selected file'}), 400 if file: filename = secure_filename(file.filename) filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) file.save(filepath) Globals.source_path = filepath return jsonify({'message': 'Source uploaded', 'filepath': filepath}), 200 @app.route('/upload/target', methods=['POST']) def upload_target(): if 'file' not in request.files: return jsonify({'error': 'No file part'}), 400 file = request.files['file'] if file.filename == '': return jsonify({'error': 'No selected file'}), 400 global target_path_web # Use the web-specific target path if file: filename = secure_filename(file.filename) filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) file.save(filepath) Globals.target_path = filepath # This is for the core processing engine target_path_web = filepath # This is for UI state, like triggering face mapping # Provide a URL to the uploaded file for preview if desired, requires a new endpoint or serving 'uploads' statically # For now, client-side preview is used. return jsonify({'message': 'Target uploaded', 'filepath': filepath, 'file_url': f'/uploads/{filename}'}), 200 @app.route('/uploads/') # Simple endpoint to serve uploaded files for preview def uploaded_file(filename): return send_from_directory(app.config['UPLOAD_FOLDER'], filename) @app.route('/update_settings', methods=['POST']) def update_settings(): data = request.get_json() if not data: return jsonify({'error': 'No data provided'}), 400 # Update Globals based on received data # Example: if 'keep_fps' in data: Globals.keep_fps = bool(data['keep_fps']) if 'keep_audio' in data: Globals.keep_audio = bool(data['keep_audio']) if 'many_faces' in data: Globals.many_faces = bool(data['many_faces']) if 'mouth_mask' in data: # HTML ID is 'mouth-mask' Globals.mouth_mask = bool(data['mouth_mask']) # Maps to Globals.mouth_mask # Add more settings as they are defined in Globals and the UI if 'frame_processors' in data: # Example for a more complex setting Globals.frame_processors = data['frame_processors'] # Assuming it's a list of strings # A more generic way if keys match Globals attributes: # for key, value in data.items(): # if hasattr(Globals, key): # # Be careful with types, e.g. ensuring booleans are booleans # if isinstance(getattr(Globals, key, None), bool): # setattr(Globals, key, bool(value)) # else: # setattr(Globals, key, value) return jsonify({'message': 'Settings updated'}), 200 @app.route('/start_processing', methods=['POST']) def start_processing(): if not Globals.source_path or not os.path.exists(Globals.source_path): return jsonify({'error': 'Source path not set or invalid'}), 400 if not Globals.target_path or not os.path.exists(Globals.target_path): return jsonify({'error': 'Target path not set or invalid'}), 400 # Determine a unique output filename and set Globals.output_path target_filename = os.path.basename(Globals.target_path) filename, ext = os.path.splitext(target_filename) unique_output_filename = f"{filename}_processed_{int(time.time())}{ext}" Globals.output_path = os.path.join(app.config['PROCESSED_OUTPUTS_FOLDER'], unique_output_filename) # Ensure default frame processors are set if none are provided by the client if not Globals.frame_processors: Globals.frame_processors = ['face_swapper'] # Default to face_swapper print("Warning: No frame processors selected by client, defaulting to 'face_swapper'.") try: # Log current settings being used print(f"Preparing to process with core engine. Source: {Globals.source_path}, Target: {Globals.target_path}, Output: {Globals.output_path}") print(f"Options: Keep FPS: {Globals.keep_fps}, Keep Audio: {Globals.keep_audio}, Many Faces: {Globals.many_faces}") print(f"Frame Processors: {Globals.frame_processors}") # Ensure necessary resources are available and limited (e.g. memory) # This was part of the old core.run() sequence. # Consider if pre_check from core should be called here too, or if it's mainly for CLI # For now, webapp assumes inputs are valid if they exist. core.limit_resources() # Call the refactored core processing function processing_result = core.process_media() if processing_result.get('success'): final_output_path = processing_result.get('output_path', Globals.output_path) # Use path from result if available # Ensure the unique_output_filename matches the actual output from process_media if it changed it # For now, we assume process_media uses Globals.output_path as set above. print(f"Core processing successful. Output at: {final_output_path}") return jsonify({ 'message': 'Processing complete', 'output_filename': os.path.basename(final_output_path), 'download_url': f'/get_output/{os.path.basename(final_output_path)}' }) else: print(f"Core processing failed: {processing_result.get('error')}") # If NSFW, include that info if process_media provides it if processing_result.get('nsfw'): return jsonify({'error': processing_result.get('error', 'NSFW content detected.'), 'nsfw': True}), 400 # Bad request due to content return jsonify({'error': processing_result.get('error', 'Unknown error during processing')}), 500 except Exception as e: # This is a fallback for unexpected errors not caught by core.process_media print(f"An unexpected error occurred in /start_processing endpoint: {e}") import traceback traceback.print_exc() return jsonify({'error': f'An critical unexpected error occurred: {str(e)}'}), 500 finally: # Always attempt to clean up temp files, regardless of success or failure # core.cleanup_temp_files() takes no args now for webapp context (quit_app=False is default) print("Executing cleanup of temporary files from webapp.") core.cleanup_temp_files() @app.route('/get_output/') def get_output(filename): return send_from_directory(app.config['PROCESSED_OUTPUTS_FOLDER'], filename, as_attachment=True) if __name__ == '__main__': # Initialize any necessary globals or configurations from core logic if needed # For example, if core.parse_args() sets up initial globals from some defaults: # import modules.core as main_core # main_core.parse_args([]) # Pass empty list or appropriate defaults if it expects CLI args # For development, directly run the Flask app. # For production, a WSGI server like Gunicorn would be used. app.run(debug=True, host='0.0.0.0', port=5000) # Video Feed Section def generate_frames(): global VIDEO_CAMERA global VIDEO_CAMERA, prev_time, frame_count, current_fps print("generate_frames: Attempting to open camera...") # Determine camera index (e.g., from Globals or default to 0) camera_index = 0 # Or Globals.camera_index if you add such a setting VIDEO_CAMERA = cv2.VideoCapture(camera_index) if not VIDEO_CAMERA.isOpened(): print(f"Error: Could not open video camera at index {camera_index}.") # TODO: Yield a placeholder image with an error message return print("generate_frames: Camera opened. Initializing settings for live processing.") prev_time = time.time() frame_count = 0 current_fps = 0 source_face = None if Globals.source_path and not Globals.map_faces: # map_faces logic for live might be complex try: source_image_cv2 = cv2.imread(Globals.source_path) if source_image_cv2 is not None: source_face = get_one_face(source_image_cv2) if source_face is None: print("Warning: No face found in source image for live preview.") except Exception as e: print(f"Error loading source image for live preview: {e}") # Get frame processors # Ensure Globals.frame_processors is a list. If it can be None, default to an empty list. current_frame_processors = Globals.frame_processors if Globals.frame_processors is not None else [] active_frame_processors = frame_processors_core.get_frame_processors_modules(current_frame_processors) # Example: Conditionally remove face enhancer if its toggle is off # This assumes fp_ui structure; adjust if it's different or not used for live mode. if not Globals.fp_ui.get('face_enhancer', False) and any(p.NAME == 'DLC.FACE-ENHANCER' for p in active_frame_processors): active_frame_processors = [p for p in active_frame_processors if p.NAME != 'DLC.FACE-ENHANCER'] print("Live Preview: Face Enhancer disabled by UI toggle.") print(f"Live Preview: Active processors: {[p.NAME for p in active_frame_processors if hasattr(p, 'NAME')]}") try: while VIDEO_CAMERA and VIDEO_CAMERA.isOpened(): # Check if VIDEO_CAMERA is not None success, frame = VIDEO_CAMERA.read() if not success: print("Error: Failed to read frame from camera during live feed.") break processed_frame = frame.copy() if Globals.live_mirror: processed_frame = cv2.flip(processed_frame, 1) # Apply Processing # Apply Processing if Globals.map_faces: if Globals.simple_map: # Check if mappings are submitted and processed for processor in active_frame_processors: if hasattr(processor, 'process_frame_v2') and callable(processor.process_frame_v2): try: processed_frame = processor.process_frame_v2(processed_frame) except Exception as e: print(f"Error applying mapped processor {processor.NAME if hasattr(processor, 'NAME') else 'Unknown'} in live feed: {e}") cv2.putText(processed_frame, "Error in mapped processing", (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) # else: No v2 method, map_faces might not apply or needs different handling else: # map_faces is true, but mappings not submitted/valid cv2.putText(processed_frame, "Map Faces: Mappings not submitted or invalid.", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) elif source_face: # Not map_faces, but single source face is available for processor in active_frame_processors: try: if hasattr(processor, 'process_frame') and callable(processor.process_frame): if processor.NAME == 'DLC.FACE-ENHANCER': processed_frame = processor.process_frame(None, processed_frame) else: processed_frame = processor.process_frame(source_face, processed_frame) except Exception as e: print(f"Error applying single source processor {processor.NAME if hasattr(processor, 'NAME') else 'Unknown'} in live feed: {e}") elif not Globals.source_path: # No map_faces and no single source image cv2.putText(processed_frame, "No Source Image Selected", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) # FPS Calculation & Overlay if Globals.show_fps: frame_count += 1 now = time.time() # Calculate FPS over a 1-second interval if (now - prev_time) > 1: current_fps = frame_count / (now - prev_time) prev_time = now frame_count = 0 cv2.putText(processed_frame, f"FPS: {current_fps:.2f}", (10, processed_frame.shape[0] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) # Encode the processed_frame to JPEG ret, buffer = cv2.imencode('.jpg', processed_frame) if not ret: print("Error: Failed to encode processed frame to JPEG.") continue frame_bytes = buffer.tobytes() yield (b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n') except GeneratorExit: print("generate_frames: Client disconnected.") except Exception as e: print(f"Exception in generate_frames main loop: {e}") import traceback traceback.print_exc() finally: print("generate_frames: Releasing camera.") if VIDEO_CAMERA: VIDEO_CAMERA.release() VIDEO_CAMERA = None # Reset global camera object @app.route('/video_feed') def video_feed(): print("Request received for /video_feed") return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame') # Optional: Endpoint to explicitly stop the camera if needed. # This is tricky with a global VIDEO_CAMERA and HTTP's stateless nature. # A more robust solution might involve websockets or a different camera management strategy. @app.route('/stop_video_feed', methods=['POST']) def stop_video_feed(): global VIDEO_CAMERA print("/stop_video_feed called") if VIDEO_CAMERA: print("Releasing video camera from /stop_video_feed") VIDEO_CAMERA.release() VIDEO_CAMERA = None return jsonify({'message': 'Video feed stopped.'}) return jsonify({'message': 'No active video feed to stop.'}) @app.route('/get_target_faces_for_mapping', methods=['GET']) def get_target_faces_for_mapping_route(): global target_path_web # Use the web-specific target path if not target_path_web or not os.path.exists(target_path_web): return jsonify({'error': 'Target image not uploaded or path is invalid.'}), 400 if not util_is_image(target_path_web): # Use the utility function for checking image type return jsonify({'error': 'Target file is not a valid image for face mapping.'}), 400 try: # This function will populate Globals.source_target_map # It expects the target image path to be in Globals.target_path for its internal logic # So, ensure Globals.target_path is also set to target_path_web for this call # This is a bit of a workaround due to how get_unique_faces_from_target_image uses Globals original_global_target_path = Globals.target_path Globals.target_path = target_path_web get_unique_faces_from_target_image() # This should fill Globals.source_target_map # Restore original Globals.target_path if it was different (e.g. from a previous full processing run) # For web UI flow, target_path_web and Globals.target_path will typically be the same after an upload. Globals.target_path = original_global_target_path if not Globals.source_target_map: return jsonify({'error': 'No faces found in the target image or error during analysis.'}), 404 response_data = [] for item in Globals.source_target_map: target_cv2_img = item['target']['cv2'] if target_cv2_img is None: # Should not happen if map is populated correctly continue _, buffer = cv2.imencode('.jpg', target_cv2_img) b64_img = base64.b64encode(buffer).decode('utf-8') response_data.append({'id': item['id'], 'image_b64': b64_img}) return jsonify(response_data) except Exception as e: print(f"Error in /get_target_faces_for_mapping: {e}") import traceback traceback.print_exc() return jsonify({'error': f'An unexpected error occurred: {str(e)}'}), 500 @app.route('/submit_face_mappings', methods=['POST']) def submit_face_mappings_route(): if 'target_ids_json' not in request.form: return jsonify({'error': 'No target_ids_json provided.'}), 400 try: target_ids = json.loads(request.form['target_ids_json']) except json.JSONDecodeError: return jsonify({'error': 'Invalid JSON in target_ids_json.'}), 400 if not Globals.source_target_map: # This implies /get_target_faces_for_mapping was not called or failed. # Or, it could be cleared. Re-populate it if target_path_web is available. if target_path_web and os.path.exists(target_path_web) and util_is_image(target_path_web): print("Re-populating source_target_map as it was empty during submit.") original_global_target_path = Globals.target_path Globals.target_path = target_path_web get_unique_faces_from_target_image() Globals.target_path = original_global_target_path if not Globals.source_target_map: return jsonify({'error': 'Could not re-initialize target faces. Please re-upload target image.'}), 500 else: return jsonify({'error': 'Target face map not initialized. Please upload target image again.'}), 500 all_mappings_valid = True processed_ids = set() for target_id_str in target_ids: target_id = int(target_id_str) # Ensure it's an integer if IDs are integers file_key = f'source_file_{target_id}' if file_key not in request.files: print(f"Warning: Source file for target_id {target_id} not found in submission.") # Mark this mapping as invalid or skip? For now, we require all submitted IDs to have files. # If a file is optional for a target, client should not include its ID in target_ids_json. # However, Globals.source_target_map will still have this target. We just won't assign a source to it. continue source_file = request.files[file_key] if source_file.filename == '': print(f"Warning: Empty filename for source file for target_id {target_id}.") continue # Skip if no file was actually selected for this input # Save the uploaded source file temporarily for this mapping temp_source_filename = f"temp_source_for_target_{target_id}_{secure_filename(source_file.filename)}" temp_source_filepath = os.path.join(app.config['UPLOAD_FOLDER'], temp_source_filename) source_file.save(temp_source_filepath) source_cv2_img = cv2.imread(temp_source_filepath) if source_cv2_img is None: print(f"Error: Could not read saved source image for target_id {target_id} from {temp_source_filepath}") # all_mappings_valid = False # Decide if one bad source fails all # os.remove(temp_source_filepath) # Clean up continue # Skip this mapping source_face_obj = get_one_face(source_cv2_img) # This also returns the cropped face usually if source_face_obj: map_entry_found = False for map_item in Globals.source_target_map: if str(map_item['id']) == str(target_id): # Compare as strings or ensure IDs are consistent type # The 'face' from get_one_face is the full Face object. # The 'cv2' image from get_one_face is the cropped face. # We need to store both, similar to how the original UI might have done. # Let's assume get_one_face returns a tuple (Face_object, cropped_cv2_image) # or that Face_object itself contains the cropped image if needed later. # For now, storing the Face object which includes embedding and bbox. # The cropped image can be re-derived or stored if `get_one_face` provides it. # Let's assume `get_one_face` is just the Face object for simplicity here, # and the cropped image for `source_target_map` needs to be handled. # A better `get_one_face` might return a dict {'face': Face, 'cv2': cropped_img} # Simplified: get_one_face returns the Face object, and we'll use that. # The `ui.update_popup_source` implies the map needs {'cv2': cropped_img, 'face': Face_obj} # Let's assume `source_face_obj` is the Face object. We need its cropped image. # This might require a helper or for get_one_face to return it. # For now, we'll store the Face object. The cropped image part for source_target_map # might need adjustment based on face_analyser's exact return for get_one_face. # A common pattern is that the Face object itself has bbox, and you can crop from original using that. # Let's assume we need to manually crop based on the Face object from get_one_face # This is a placeholder - exact cropping depends on what get_one_face returns and what processors need # For now, we'll just store the Face object. # If `face_swapper`'s `process_frame_v2` needs cropped source images in `source_target_map`, # this part needs to ensure they are correctly populated. # For simplicity, assuming `get_one_face` returns the main `Face` object, and `face_swapper` can use that. # The `source_target_map` structure is critical. # Looking at `face_swapper.py`, `process_frame_v2` uses `Globals.simple_map`. # `simplify_maps()` populates `simple_map` from `source_target_map`. # `simplify_maps()` expects `item['source']['face']` to be the source `Face` object. map_item['source'] = {'face': source_face_obj, 'cv2': source_cv2_img} # Store the original uploaded source, not necessarily cropped yet. Processors handle cropping. map_entry_found = True processed_ids.add(target_id) break if not map_entry_found: print(f"Warning: Target ID {target_id} from submission not found in existing map.") all_mappings_valid = False # Or handle as error else: print(f"Warning: No face found in uploaded source for target_id {target_id}.") # Mark this specific mapping as invalid by not adding a 'source' to it, or removing it. # For now, we just don't add a source. simplify_maps should handle items without a source. all_mappings_valid = False # if strict, one failed source makes all invalid for this submission batch # Clean up the temporary saved source file if os.path.exists(temp_source_filepath): os.remove(temp_source_filepath) # Clear 'source' for any target_ids that were in source_target_map but not in this submission # or if their source file didn't yield a face. for map_item in Globals.source_target_map: if map_item['id'] not in processed_ids and 'source' in map_item: del map_item['source'] if not all_mappings_valid: # Or based on a stricter check # simplify_maps() will still run and create mappings for valid pairs print("simplify_maps: Some mappings may be invalid or incomplete.") simplify_maps() # Populate Globals.simple_map based on updated Globals.source_target_map # For debugging: # print("Updated source_target_map:", Globals.source_target_map) # print("Generated simple_map:", Globals.simple_map) if not Globals.simple_map and all_mappings_valid and target_ids: # If all submitted were meant to be valid but simple_map is empty return jsonify({'error': 'Mappings processed, but no valid face pairs were established. Check source images.'}), 400 Globals.map_faces = True # Crucial: Set this global so processing functions know to use the map return jsonify({'message': 'Face mappings submitted and processed.'}) # except Exception as e: # print(f"Error in /submit_face_mappings: {e}") # import traceback # traceback.print_exc() # return jsonify({'error': f'An unexpected error occurred: {str(e)}'}), 500