ai-flow-python/app/service/task/impl/smart_beauty_task.py

import logging
import os
import subprocess
import time
from typing import Any

from app.service.task.base_task import BaseTask
from app.service.task.task_param import TaskParam
from app.service.task.task_result import TaskResult
from app.util.res_util import get_resources_path

logger = logging.getLogger(__name__)


class SmartBeautyTask(BaseTask):
    def get_type(self) -> str:
        return "smart_beauty"

    def execute(self, param: TaskParam) -> TaskResult:
        try:
            images_obj = param.inputs.get("images")
            images_in: list[str] = images_obj if isinstance(images_obj, list) else []

            if not images_in:
                return TaskResult.failure("No images provided for processing")

            workspace = param.workspace
            resources = param.resources
            if not workspace or not os.path.isdir(workspace):
                return TaskResult.failure("Workspace path is not provided")

            logger.info("[smart_beauty] Starting: %d images, resources=%s", len(images_in), resources)
            t0 = time.time()
            images_out: list[str] = []
            for i, inp in enumerate(images_in):
                try:
                    base_name = _get_base_name(inp)
                    output_path = f"{param.step_id}/{base_name}_beauty.jpg"
                    logger.info("[smart_beauty] Processing image %d/%d: %s -> %s",
                                 i + 1, len(images_in), inp, output_path)
                    if self._execute_beauty(workspace, inp, output_path, resources):
                        images_out.append(output_path)
                        logger.info("[smart_beauty] Image %d/%d processed successfully", i + 1, len(images_in))
                    else:
                        logger.warning("[smart_beauty] Image %d/%d processing failed", i + 1, len(images_in))
                except Exception as e:
                    logger.error("[smart_beauty] Error processing image %d/%d %s: %s",
                                  i + 1, len(images_in), inp, e)

            if not images_out:
                return TaskResult.failure("No images were successfully processed")

            logger.info("[smart_beauty] Completed: %d/%d images in %.2fs",
                         len(images_out), len(images_in), time.time() - t0)
            return TaskResult.success({"images": images_out})
        except Exception as e:
            logger.error("[smart_beauty] Error: %s", e, exc_info=True)
            return TaskResult.failure(f"Smart beauty processing failed: {e}")

    @staticmethod
    def _execute_beauty(workspace: str, inp: str, output: str, config: str) -> bool:
        try:
            base_dir = os.path.join(get_resources_path(), "beautify")
            executor_path = os.path.join(base_dir, "beautify")
            lib_path = os.path.join(base_dir, "libs")
            licence_path = os.path.join(base_dir, "res", "license.lic")
            face_model = os.path.join(base_dir, "res", "Face_Video_Template_p_4.0.0.model")
            hair_model = os.path.join(base_dir, "res", "Segment_Hair_p_4.4.0.model")

            if not os.path.isfile(executor_path):
                logger.error("[smart_beauty] Beauty program not found: %s", executor_path)
                return False

            cmd = [
                executor_path,
                f"-i={inp}",
                f"-o={output}",
                f"-p={config}",
                f"-l={licence_path}",
                f"-face={face_model}",
                f"-hair={hair_model}",
            ]

            env = os.environ.copy()
            ld = env.get("LD_LIBRARY_PATH", "")
            env["LD_LIBRARY_PATH"] = f"{ld}:{lib_path}" if ld else lib_path

            logger.info("[smart_beauty] Executing: %s", " ".join(cmd))
            cmd_t0 = time.time()
            result = subprocess.run(
                cmd, cwd=workspace, env=env,
                capture_output=True, text=True, timeout=60,
            )
            cmd_elapsed = time.time() - cmd_t0

            if result.stdout:
                logger.info("[smart_beauty] stdout: %s", result.stdout.strip())
            if result.stderr:
                logger.warning("[smart_beauty] stderr: %s", result.stderr.strip())

            if result.returncode == 0:
                out_file = os.path.join(workspace, output)
                if os.path.isfile(out_file):
                    file_size = os.path.getsize(out_file)
                    logger.info("[smart_beauty] Output generated: %s (%.1f KB) in %.2fs",
                                 output, file_size / 1024, cmd_elapsed)
                    return True
                logger.warning("[smart_beauty] Output file not generated: %s", output)
                return False
            else:
                logger.error("[smart_beauty] Process failed: exit_code=%d in %.2fs",
                              result.returncode, cmd_elapsed)
                return False
        except subprocess.TimeoutExpired:
            logger.error("[smart_beauty] Process timeout (60s) for input: %s", inp)
            return False
        except Exception as e:
            logger.error("[smart_beauty] Error executing command: %s", e)
            return False


def _get_base_name(path: str) -> str:
    name = os.path.basename(path)
    dot = name.rfind(".")
    return name[:dot] if dot > 0 else name