src/main.py

import logging
from pathlib import Path
from typing import Union
from argparse import ArgumentParser
import cv2
from PIL import Image

# from src.analyze import analyze_and_visualize
from src.skin_analyzer import analyze_skin_function
from src.image import ImageBundle
import os
import numpy as np
import timeit

LOG = logging.getLogger(__name__)


def process_image(
    filename_or_url: Union[str, Path], analyze_skin, analyze_eyes, analyze_hair
):
    """
    Process the image and return the result.
    :param filename_or_url: The filename (in local devices) or URL (in Internet) of the image.
    :param analyze_skin: Whether to perform skin analysis.
    :param analyze_eyes: Whether to perform eye analysis.
    :param analyze_hair: Whether to perform hair analysis.
    :return: Analysis results.
    """
    # Load the image
    # time it
    start = timeit.default_timer()
    image_bundle = ImageBundle(image_source=filename_or_url)
    stop = timeit.default_timer()
    print("load image time: ", stop - start)

    # Detect faces and landmarks
    start = timeit.default_timer()
    face_data = image_bundle.detect_faces_and_landmarks()
    stop = timeit.default_timer()
    print("detect faces and landmarks time: ", stop - start)
    start = timeit.default_timer()
    landmarks = image_bundle.detect_face_landmarks()
    stop = timeit.default_timer()
    print("detect face landmarks time: ", stop - start)

    start = timeit.default_timer()
    # Perform segmentation
    segmentation_maps = image_bundle.segment_image()
    stop = timeit.default_timer()
    print("segment image time: ", stop - start)

    analysis_results = {}
    overlay_images = []

    image_np = image_bundle.numpy_image()
    print("image_np shape: ", image_np.shape)
    # import IPython; IPython.embed()

    if analyze_hair:
        hair_mask = segmentation_maps["hair_mask"]
        hair_analysis = analyze_hair_function(image_np, hair_mask)
        analysis_results["hair_analysis"] = hair_analysis
        overlay_images.append(segmentation_maps["hair_mask"])

    if analyze_skin:
        skin_mask = segmentation_maps["face_skin_mask"]

        skin_analysis = analyze_skin_function(image_np, skin_mask)
        overlay_images.append(skin_analysis["overlay_image"])
        if "overlay_image" in skin_analysis.keys():
            del skin_analysis["overlay_image"]
        if "filtered_skin_mask" in skin_analysis.keys():
            del skin_analysis["filtered_skin_mask"]
        if "dominant_skin_tones" in skin_analysis.keys():
            del skin_analysis["dominant_skin_tones"]
        analysis_results["skin_analysis"] = skin_analysis

        # del overlay_images[-1]

    if analyze_eyes:
        eye_mask = (
            segmentation_maps["right_eye_mask"] | segmentation_maps["left_eye_mask"]
        )
        eye_analysis = analyze_eye_function(image_np, eye_mask)
        analysis_results["eye_analysis"] = eye_analysis
        overlay_images.append(segmentation_maps["right_eye_mask"])

    # Combine overlay images
    if overlay_images:
        combined_overlay = np.maximum.reduce(overlay_images)
    else:
        combined_overlay = image_np

    # Convert combined_overlay to BGR for saving with OpenCV
    combined_overlay_bgr = cv2.cvtColor(combined_overlay, cv2.COLOR_RGB2BGR)

    return combined_overlay_bgr, analysis_results


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-i", "--image_path", required=True, help="Path to the image")
    parser.add_argument(
        "--analyze_skin", action="store_true", help="Perform skin analysis"
    )
    parser.add_argument(
        "--analyze_eyes", action="store_true", help="Perform eye analysis"
    )
    parser.add_argument(
        "--analyze_hair", action="store_true", help="Perform hair analysis"
    )

    args = parser.parse_args()
    image_path = args.image_path

    # Process the image
    combined_overlay, analysis_results = process_image(
        image_path,
        analyze_skin=args.analyze_skin,
        analyze_eyes=args.analyze_eyes,
        analyze_hair=args.analyze_hair,
    )

    # Save the results
    im_basename = os.path.basename(image_path).split(".")[0]
    overlay_image_path = f"outputs/overlay_{im_basename}.png"
    cv2.imwrite(overlay_image_path, combined_overlay)
    import json

    # Save the analysis results as JSON, including the path to the overlay image
    analysis_results["overlay_image_path"] = overlay_image_path
    print(analysis_results)

    with open(f"outputs/analysis_{im_basename}.json", "w") as f:
        json.dump(analysis_results, f, indent=4)

    print("Analysis complete. Results saved in 'outputs' directory.")