Clean up main.py

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

304
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 multiprocessing
import os
import threading
import uuid
from dataclasses import dataclass
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
API_KEY = os.environ.get("IMMICH_API_KEY", "")
BASE_URL = os.environ.get("IMMICH_BASE_URL", "")
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
# Configure logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# Filter out progress route logs
class ProgressRouteFilter(logging.Filter):
def filter(self, record):
# Filter out logs containing the progress route
def filter(self, record: logging.LogRecord) -> bool:
return "/progress" not in record.getMessage()
log = logging.getLogger('werkzeug')
log.addFilter(ProgressRouteFilter())
def update_progress(current, total):
"""
Updates the global progress dictionary.
@dataclass
class AppConfig:
"""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:
current (int): Number of completed tasks.
total (int): Total number of tasks.
current: Number of completed tasks
total: Total number of tasks
"""
progress_info["completed"] = current
progress_info["total"] = total
progress_info["status"] = "running" if current < total else "done"
def background_process(person_id, resize_size, face_resolution_threshold, pose_threshold,
left_eye_pos, output_folder, api_key, base_url, progress_callback=None, cancel_flag=None):
def check_output_folder() -> Tuple[bool, int]:
"""Check if the output folder is empty.
Returns:
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:
person_id (str): ID of the person to process.
resize_size (int): Size to resize the output images to.
face_resolution_threshold (int): Minimum face resolution threshold.
pose_threshold (float): Maximum allowed head pose deviation.
left_eye_pos (tuple): Desired position of the left eye in the output.
output_folder (str): Folder to save the output images.
api_key (str): API key for authentication.
base_url (str): Base URL of the API.
progress_callback (callable, optional): Callback for progress updates.
cancel_flag (callable, optional): Function to check if process should be cancelled.
person_id: ID of the person to process
resize_size: Size to resize output images to
face_resolution_threshold: Minimum face resolution threshold
pose_threshold: Maximum allowed head pose deviation
left_eye_pos: Desired position of the left eye in output
output_folder: Folder to save output images
date_from: Optional start date in YYYY-MM-DD format
date_to: Optional end date in YYYY-MM-DD format
progress_callback: Optional callback for progress updates
cancel_flag: Optional function to check for cancellation
Returns:
List of processed file paths
"""
try:
config = ProcessConfig(
api_key=api_key,
base_url=base_url,
process_config = ProcessConfig(
api_key=config.api_key,
base_url=config.base_url,
person_id=person_id,
output_folder=output_folder,
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_height=face_resolution_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
processed_files = process_faces(
config=config,
return process_faces(
config=process_config,
max_workers=1,
progress_callback=progress_callback,
cancel_flag=cancel_flag
)
return processed_files
except Exception as e:
logger.error(f"Error in background process: {str(e)}")
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")
def progress():
"""
Endpoint to return current progress as JSON.
"""
def progress() -> Dict[str, any]:
"""Get current progress information."""
return jsonify(progress_info)
@app.route("/check-connection")
def check_connection():
"""
Endpoint to check the Immich server connection.
"""
is_valid, message = validate_immich_connection(API_KEY, BASE_URL)
def check_connection() -> Dict[str, any]:
"""Check connection to Immich server."""
is_valid, message = validate_immich_connection(config.api_key, config.base_url)
return jsonify({"valid": is_valid, "message": message})
@app.route("/cancel", methods=["POST"])
def cancel():
"""
Endpoint to cancel the current processing job.
"""
def cancel() -> Dict[str, any]:
"""Cancel the current processing job."""
global processing_thread, cancel_requested
# Set the cancel flag
cancel_requested = True
if processing_thread and processing_thread.is_alive():
progress_info["status"] = "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"])
def index():
"""
Index route that displays the form and starts processing in a background thread on POST.
"""
def index() -> str:
"""Handle the main page and processing requests."""
global processing_thread, cancel_requested
result = None
error = None
warning = None
# Check if output folder is empty
# Check output folder status
is_empty, file_count = check_output_folder()
if not is_empty:
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
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)))
# Validate connection on 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:
# Make sure any previous cancel request is cleared
cancel_requested = False
# Get form data with defaults
person_id = request.form["person_id"]
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)
resize_size = int(request.form.get("resize_size", config.default_resize_size))
face_resolution_threshold = int(request.form.get("face_resolution_threshold",
config.default_face_resolution_threshold))
pose_threshold = float(request.form.get("pose_threshold", config.default_pose_threshold))
# Date ranges are optional
# Optional date ranges
date_from = request.form.get("date_from") or None
date_to = request.form.get("date_to") or None
# Video compilation options
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
progress_info["completed"] = 0
progress_info["total"] = 0
progress_info["status"] = "idle"
# Reset progress info
progress_info.update({
"completed": 0,
"total": 0,
"status": "idle"
})
progress_info.pop("video_path", None)
# Start the processing in a background thread
# Start processing
processing_thread = threading.Thread(
target=background_process,
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()
result = "Processing started. Please wait and watch the progress bar below."
except Exception as e:
error = f"Error processing request: {e}"
return render_template("index.html", result=result, error=error, warning=warning,
max_workers_options=list(range(1, AVAILABLE_CORES + 1)))
return render_template("index.html",
result=result,
error=error,
warning=warning,
max_workers_options=list(range(1, AVAILABLE_CORES + 1)))
if __name__ == "__main__":
app.run(host="0.0.0.0", port=5000)
app.run(debug=True)