Clean up main.py

This commit is contained in:
Arnaud_Cayrol 2025-04-12 13:54:27 +02:00
parent fc856e10c3
commit b5d340c1ae

300
main.py
View file

@ -1,74 +1,116 @@
import os
import multiprocessing
import threading
import subprocess
from flask import Flask, request, render_template, jsonify, redirect, url_for
from timelapse import process_faces, ProcessConfig, validate_immich_connection
import logging import logging
import multiprocessing
import os
import threading
import uuid import uuid
from dataclasses import dataclass
from datetime import datetime from datetime import datetime
from typing import Callable, Dict, List, Optional, Tuple
app = Flask(__name__) from flask import Flask, jsonify, render_template, request
from timelapse import ProcessConfig, process_faces, validate_immich_connection
# Critical parameters provided via environment variables # Configure logging
API_KEY = os.environ.get("IMMICH_API_KEY", "") logger = logging.getLogger(__name__)
BASE_URL = os.environ.get("IMMICH_BASE_URL", "") logger.setLevel(logging.INFO)
OUTPUT_FOLDER = "output"
# Model paths
LANDMARK_MODEL = "shape_predictor_68_face_landmarks.dat"
AVAILABLE_CORES = multiprocessing.cpu_count()
# Global progress dictionary only one job at a time is assumed here
progress_info = {"completed": 0, "total": 0, "status": "idle"}
# Global processing thread reference
processing_thread = None
# Global flag to signal cancellation
cancel_requested = False
# Filter out progress route logs
class ProgressRouteFilter(logging.Filter): class ProgressRouteFilter(logging.Filter):
def filter(self, record): def filter(self, record: logging.LogRecord) -> bool:
# Filter out logs containing the progress route
return "/progress" not in record.getMessage() return "/progress" not in record.getMessage()
log = logging.getLogger('werkzeug') log = logging.getLogger('werkzeug')
log.addFilter(ProgressRouteFilter()) log.addFilter(ProgressRouteFilter())
def update_progress(current, total): @dataclass
""" class AppConfig:
Updates the global progress dictionary. """Configuration for the application."""
api_key: str
base_url: str
output_folder: str
landmark_model: str
default_resize_size: int = 512
default_face_resolution_threshold: int = 128
default_pose_threshold: float = 25.0
default_left_eye_pos: Tuple[float, float] = (0.35, 0.4)
default_framerate: int = 24
default_date_format: str = "%Y-%m-%d"
# Initialize configuration
config = AppConfig(
api_key="wHNgNvlsWiSqUWvjR2K8kK2ngToQXlMIiGwF5LkxY",
base_url="http://192.168.1.94:2283/api",
output_folder="output",
landmark_model="shape_predictor_68_face_landmarks.dat"
)
# Initialize Flask app
app = Flask(__name__)
# Global state
AVAILABLE_CORES = multiprocessing.cpu_count()
progress_info: Dict[str, any] = {"completed": 0, "total": 0, "status": "idle"}
processing_thread: Optional[threading.Thread] = None
cancel_requested: bool = False
def update_progress(current: int, total: int) -> None:
"""Update the global progress information.
Args: Args:
current (int): Number of completed tasks. current: Number of completed tasks
total (int): Total number of tasks. total: Total number of tasks
""" """
progress_info["completed"] = current progress_info["completed"] = current
progress_info["total"] = total progress_info["total"] = total
progress_info["status"] = "running" if current < total else "done" progress_info["status"] = "running" if current < total else "done"
def check_output_folder() -> Tuple[bool, int]:
"""Check if the output folder is empty.
def background_process(person_id, resize_size, face_resolution_threshold, pose_threshold, Returns:
left_eye_pos, output_folder, api_key, base_url, progress_callback=None, cancel_flag=None): Tuple containing (is_empty, file_count)
""" """
Background process to handle face alignment and timelapse creation. if not os.path.exists(config.output_folder):
os.makedirs(config.output_folder, exist_ok=True)
return True, 0
files = [f for f in os.listdir(config.output_folder)
if os.path.isfile(os.path.join(config.output_folder, f))]
return len(files) == 0, len(files)
def background_process(
person_id: str,
resize_size: int,
face_resolution_threshold: int,
pose_threshold: float,
left_eye_pos: Tuple[float, float],
output_folder: str,
date_from: Optional[str] = None,
date_to: Optional[str] = None,
progress_callback: Optional[Callable] = None,
cancel_flag: Optional[Callable] = None
) -> List[str]:
"""Process faces in the background.
Args: Args:
person_id (str): ID of the person to process. person_id: ID of the person to process
resize_size (int): Size to resize the output images to. resize_size: Size to resize output images to
face_resolution_threshold (int): Minimum face resolution threshold. face_resolution_threshold: Minimum face resolution threshold
pose_threshold (float): Maximum allowed head pose deviation. pose_threshold: Maximum allowed head pose deviation
left_eye_pos (tuple): Desired position of the left eye in the output. left_eye_pos: Desired position of the left eye in output
output_folder (str): Folder to save the output images. output_folder: Folder to save output images
api_key (str): API key for authentication. date_from: Optional start date in YYYY-MM-DD format
base_url (str): Base URL of the API. date_to: Optional end date in YYYY-MM-DD format
progress_callback (callable, optional): Callback for progress updates. progress_callback: Optional callback for progress updates
cancel_flag (callable, optional): Function to check if process should be cancelled. cancel_flag: Optional function to check for cancellation
Returns:
List of processed file paths
""" """
try: try:
config = ProcessConfig( process_config = ProcessConfig(
api_key=api_key, api_key=config.api_key,
base_url=base_url, base_url=config.base_url,
person_id=person_id, person_id=person_id,
output_folder=output_folder, output_folder=output_folder,
resize_width=resize_size, resize_width=resize_size,
@ -76,182 +118,114 @@ def background_process(person_id, resize_size, face_resolution_threshold, pose_t
min_face_width=face_resolution_threshold, min_face_width=face_resolution_threshold,
min_face_height=face_resolution_threshold, min_face_height=face_resolution_threshold,
pose_threshold=pose_threshold, pose_threshold=pose_threshold,
left_eye_pos=left_eye_pos left_eye_pos=left_eye_pos,
landmark_model_path=config.landmark_model,
date_from=date_from,
date_to=date_to
) )
# Process the faces return process_faces(
processed_files = process_faces( config=process_config,
config=config,
max_workers=1, max_workers=1,
progress_callback=progress_callback, progress_callback=progress_callback,
cancel_flag=cancel_flag cancel_flag=cancel_flag
) )
return processed_files
except Exception as e: except Exception as e:
logger.error(f"Error in background process: {str(e)}") logger.error(f"Error in background process: {str(e)}")
raise raise
def check_output_folder():
"""
Checks if the output folder is empty.
Returns:
tuple: (is_empty, file_count) - Boolean indicating if folder is empty and number of files
"""
if not os.path.exists(OUTPUT_FOLDER):
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
return True, 0
files = [f for f in os.listdir(OUTPUT_FOLDER) if os.path.isfile(os.path.join(OUTPUT_FOLDER, f))]
return len(files) == 0, len(files)
@app.route("/progress") @app.route("/progress")
def progress(): def progress() -> Dict[str, any]:
""" """Get current progress information."""
Endpoint to return current progress as JSON.
"""
return jsonify(progress_info) return jsonify(progress_info)
@app.route("/check-connection") @app.route("/check-connection")
def check_connection(): def check_connection() -> Dict[str, any]:
""" """Check connection to Immich server."""
Endpoint to check the Immich server connection. is_valid, message = validate_immich_connection(config.api_key, config.base_url)
"""
is_valid, message = validate_immich_connection(API_KEY, BASE_URL)
return jsonify({"valid": is_valid, "message": message}) return jsonify({"valid": is_valid, "message": message})
@app.route("/cancel", methods=["POST"]) @app.route("/cancel", methods=["POST"])
def cancel(): def cancel() -> Dict[str, any]:
""" """Cancel the current processing job."""
Endpoint to cancel the current processing job.
"""
global processing_thread, cancel_requested global processing_thread, cancel_requested
# Set the cancel flag
cancel_requested = True cancel_requested = True
if processing_thread and processing_thread.is_alive(): if processing_thread and processing_thread.is_alive():
progress_info["status"] = "cancelled" progress_info["status"] = "cancelled"
return jsonify({"success": True, "message": "Processing cancelled."}) return jsonify({"success": True, "message": "Processing cancelled."})
else:
cancel_requested = False
return jsonify({"success": False, "message": "No active processing to cancel."})
@app.route("/process", methods=["POST"])
def process():
"""Handle the processing request."""
try:
person_id = request.form.get("person_id")
if not person_id:
return jsonify({"error": "Person ID is required"}), 400
resize_size = int(request.form.get("resize_size", 512))
face_resolution_threshold = int(request.form.get("face_resolution_threshold", 128))
pose_threshold = float(request.form.get("pose_threshold", 25))
left_eye_x = float(request.form.get("left_eye_x", 0.4))
left_eye_y = float(request.form.get("left_eye_y", 0.4))
left_eye_pos = (left_eye_x, left_eye_y)
# Create output folder with timestamp
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_folder = os.path.join("output", f"timelapse_{timestamp}")
os.makedirs(output_folder, exist_ok=True)
# Start background process
process_id = str(uuid.uuid4())
process_info = {
"status": "running",
"start_time": datetime.now().isoformat(),
"output_folder": output_folder
}
active_processes[process_id] = process_info
# Start the background process
process = multiprocessing.Process(
target=background_process,
args=(person_id, resize_size, face_resolution_threshold, pose_threshold,
left_eye_pos, output_folder, API_KEY, BASE_URL)
)
process.start()
active_processes[process_id]["process"] = process
return jsonify({"process_id": process_id})
except Exception as e:
logger.error(f"Error in process route: {str(e)}")
return jsonify({"error": str(e)}), 500
cancel_requested = False
return jsonify({"success": False, "message": "No active processing to cancel."})
@app.route("/", methods=["GET", "POST"]) @app.route("/", methods=["GET", "POST"])
def index(): def index() -> str:
""" """Handle the main page and processing requests."""
Index route that displays the form and starts processing in a background thread on POST.
"""
global processing_thread, cancel_requested global processing_thread, cancel_requested
result = None result = None
error = None error = None
warning = None warning = None
# Check if output folder is empty # Check output folder status
is_empty, file_count = check_output_folder() is_empty, file_count = check_output_folder()
if not is_empty: if not is_empty:
warning = f"Output folder is not empty. Contains {file_count} files. New images will be added to this folder." warning = f"Output folder is not empty. Contains {file_count} files. New images will be added to this folder."
# Check if Immich server connection is valid # Validate connection on POST
is_valid, message = validate_immich_connection(API_KEY, BASE_URL)
if not is_valid and request.method == "POST":
error = f"Immich server connection error: {message}"
return render_template("index.html", error=error, warning=warning,
max_workers_options=list(range(1, AVAILABLE_CORES + 1)))
if request.method == "POST": if request.method == "POST":
is_valid, message = validate_immich_connection(config.api_key, config.base_url)
if not is_valid:
error = f"Immich server connection error: {message}"
return render_template("index.html", error=error, warning=warning,
max_workers_options=list(range(1, AVAILABLE_CORES + 1)))
try: try:
# Make sure any previous cancel request is cleared
cancel_requested = False cancel_requested = False
# Get form data with defaults
person_id = request.form["person_id"] person_id = request.form["person_id"]
resize_size = int(request.form.get("resize_size", 512)) resize_size = int(request.form.get("resize_size", config.default_resize_size))
face_resolution_threshold = int(request.form.get("face_resolution_threshold", 128)) face_resolution_threshold = int(request.form.get("face_resolution_threshold",
pose_threshold = float(request.form.get("pose_threshold", 25)) config.default_face_resolution_threshold))
left_eye_x = float(request.form.get("left_eye_x", 0.4)) pose_threshold = float(request.form.get("pose_threshold", config.default_pose_threshold))
left_eye_y = float(request.form.get("left_eye_y", 0.4))
left_eye_pos = (left_eye_x, left_eye_y)
# Date ranges are optional # Optional date ranges
date_from = request.form.get("date_from") or None date_from = request.form.get("date_from") or None
date_to = request.form.get("date_to") or None date_to = request.form.get("date_to") or None
# Video compilation options
compile_video = request.form.get("compile_video") == "on" compile_video = request.form.get("compile_video") == "on"
framerate = int(request.form.get("framerate", 24)) framerate = int(request.form.get("framerate", config.default_framerate))
# Reset progress info before starting # Reset progress info
progress_info["completed"] = 0 progress_info.update({
progress_info["total"] = 0 "completed": 0,
progress_info["status"] = "idle" "total": 0,
"status": "idle"
})
progress_info.pop("video_path", None) progress_info.pop("video_path", None)
# Start the processing in a background thread # Start processing
processing_thread = threading.Thread( processing_thread = threading.Thread(
target=background_process, target=background_process,
args=(person_id, resize_size, face_resolution_threshold, pose_threshold, args=(person_id, resize_size, face_resolution_threshold, pose_threshold,
left_eye_pos, OUTPUT_FOLDER, API_KEY, BASE_URL, update_progress, lambda: cancel_requested) config.default_left_eye_pos, config.output_folder, date_from, date_to,
update_progress, lambda: cancel_requested)
) )
processing_thread.start() processing_thread.start()
result = "Processing started. Please wait and watch the progress bar below." result = "Processing started. Please wait and watch the progress bar below."
except Exception as e: except Exception as e:
error = f"Error processing request: {e}" error = f"Error processing request: {e}"
return render_template("index.html", result=result, error=error, warning=warning, return render_template("index.html",
max_workers_options=list(range(1, AVAILABLE_CORES + 1))) result=result,
error=error,
warning=warning,
max_workers_options=list(range(1, AVAILABLE_CORES + 1)))
if __name__ == "__main__": if __name__ == "__main__":
app.run(host="0.0.0.0", port=5000) app.run(debug=True)