app.py

import gradio as gr
import cv2
import numpy as np
from PIL import Image

from src.skin_analyzer import analyze_skin_function
from src.image import ImageBundle
import plotly.express as px
import plotly.graph_objects as go

import matplotlib.pyplot as plt
import io


def resize_image(image, height=512):
    # Calculate the new width to maintain the aspect ratio
    aspect_ratio = image.width / image.height
    new_width = int(height * aspect_ratio)
    return image.resize((new_width, height), Image.Resampling.LANCZOS)


def create_histogram(data, title):
    fig = px.histogram(data, nbins=30, title=title)
    fig.update_layout(title=dict(y=0.9))
    fig.update_xaxes(title_text="Value")
    fig.update_yaxes(title_text="Frequency")
    return fig


# Function to create pie chart using Plotly
def create_pie_chart(data, title):
    fig = px.pie(values=data.values(), names=data.keys(), title=title)
    fig.update_layout(title=dict(y=0.9))
    return fig


# Function to create bar chart using Plotly
def create_bar_chart(data, title, top_n=None):
    sorted_data = dict(sorted(data.items(), key=lambda item: item[1], reverse=True))
    colors = (
        ["red" if i < top_n else "blue" for i in range(len(sorted_data))]
        if top_n and top_n < len(sorted_data)
        else "blue"
    )
    fig = px.bar(x=list(sorted_data.keys()), y=list(sorted_data.values()), title=title)
    fig.update_traces(marker_color=colors)
    fig.update_layout(
        title=dict(y=0.9), xaxis_title=None, yaxis_title="Counts", showlegend=False
    )
    fig.update_xaxes(tickangle=45)
    return fig


def process_image(
    image,
    l_min_skin,
    l_max_skin,
    l_min_tonality,
    l_max_tonality,
    chroma_thresh,
    analyze_skin,
    analyze_eyes,
    analyze_hair,
):
    """
    Process the uploaded image and return the analysis results based on selected analyses.
    :param image: Uploaded 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.
    """
    image_bundle = ImageBundle(image_array=image)

    # Detect faces and landmarks
    face_data = image_bundle.detect_faces_and_landmarks()
    landmarks = image_bundle.detect_face_landmarks()

    # Perform segmentation
    segmentation_maps = image_bundle.segment_image()

    analysis_results = {}
    # overlay_images = []

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

    if analyze_skin:
        skin_mask = segmentation_maps["face_skin_mask"]
        skin_analysis = analyze_skin_function(
            image,
            skin_mask,
            l_min_skin,
            l_max_skin,
            l_min_tonality,
            l_max_tonality,
            chroma_thresh,
        )
        if "filtered_skin_mask" in skin_analysis:
            filtered_skin_mask = skin_analysis["filtered_skin_mask"]
            del skin_analysis["filtered_skin_mask"]
        if "l_hist" in skin_analysis:
            l_hist = skin_analysis["l_hist"]
            del skin_analysis["l_hist"]
        if "tonality_l_hist" in skin_analysis:
            tonality_l_hist = skin_analysis["tonality_l_hist"]
            del skin_analysis["tonality_l_hist"]
        if "chroma_hist" in skin_analysis:
            chroma_hist = skin_analysis["chroma_hist"]
            del skin_analysis["chroma_hist"]
        # analysis_results["skin_analysis"] = skin_analysis
        # Create bar charts for analysis results
        chroma_chart = create_pie_chart(skin_analysis["chroma_counts"], "Chroma Counts")
        undertone_chart = create_pie_chart(
            skin_analysis["undertone_counts"], "Undertone Counts"
        )
        overtone_chart = create_bar_chart(
            skin_analysis["overtone_counts"], "Overtone Counts", 1
        )
        tonality_chart = create_pie_chart(
            skin_analysis["tonality_counts"], "Tonality Counts"
        )
        season_chart = create_bar_chart(
            skin_analysis["season_counts"], "Season Counts", 3
        )

        # overlay_images.append(skin_analysis["overlay_image"])

    # if analyze_eyes:
    #     eye_mask = segmentation_maps["right_iris_mask"] | segmentation_maps["left_iris_mask"]
    #     # analysis_results['eye_analysis'] = analyze_eye_function(image, eye_mask)
    #     overlay_images.append(eye_mask)

    # Combine overlay images
    overlay = image.copy()
    overlay[filtered_skin_mask > 0] = (0, 0, 255)  # Red for skin

    overlay = cv2.addWeighted(image, 0.85, overlay, 0.15, 0)

    # Convert combined_overlay to PIL Image for display
    combined_overlay = Image.fromarray(overlay)
    # Resize images before returning
    combined_overlay = resize_image(combined_overlay)
    # l_hist = resize_image(l_hist)
    # tonality_l_hist = resize_image(tonality_l_hist)
    # chroma_hist = resize_image(chroma_hist)
    # chroma_chart = resize_image(chroma_chart)
    # undertone_chart = resize_image(undertone_chart)
    # overtone_chart = resize_image(overtone_chart)
    # tonality_chart = resize_image(tonality_chart)
    # season_chart = resize_image(season_chart)

    return (
        combined_overlay,
        l_hist,
        tonality_l_hist,
        chroma_hist,
        # analysis_results,
        chroma_chart,
        undertone_chart,
        overtone_chart,
        tonality_chart,
        season_chart,
    )


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            skin_checkbox = gr.Checkbox(label="Skin Analysis", value=True)
            submit_button = gr.Button("Submit")
        with gr.Column():
            eye_checkbox = gr.Checkbox(label="Eye Analysis", value=False)
            hair_checkbox = gr.Checkbox(label="Hair Analysis", value=False)
    with gr.Row():
        with gr.Column():
            upload_image = gr.Image(type="numpy", label="Upload an Image")
        with gr.Column():
            l_min_slider = gr.Slider(
                minimum=0, maximum=100, value=10, label="L(%) Min Skin"
            )
            l_max_slider = gr.Slider(
                minimum=0, maximum=100, value=90, label="L(%) Max Skin"
            )
            l_hist_output = gr.Plot(label="L Histogram")

            # processed_image_output = gr.Image(type="pil", label="Processed Image")
    with gr.Row():
        with gr.Column():
            processed_image_output = gr.Image(type="pil", label="Processed Image")
        with gr.Column():
            undertone_chart_output = gr.Plot(label="Undertone Counts")
    with gr.Row():
        with gr.Column():
            tonality_min_slider = gr.Slider(
                minimum=0, maximum=100, value=50, label="L(%) Min Tonality"
            )
            tonality_max_slider = gr.Slider(
                minimum=0, maximum=100, value=70, label="L(%) Max Tonality"
            )
            tonality_hist_output = gr.Plot(label="Tonality L Histogram")
        with gr.Column():
            tonality_chart_output = gr.Plot(label="Tonality Counts")

    # with gr.Row():
    #     submit_button = gr.Button("Submit")

    with gr.Row():
        with gr.Column():
            chroma_slider = gr.Slider(
                minimum=0, maximum=100, value=50, label="Chroma(%) Threshold"
            )
            chroma_hist_output = gr.Plot(label="Chroma Histogram")
        with gr.Column():
            chroma_chart_output = gr.Plot(label="Chroma Counts")

    with gr.Row():
        with gr.Column():
            overtone_chart_output = gr.Plot(label="Overtone Counts")
        with gr.Column():
            season_chart_output = gr.Plot(label="Season Counts")

    inputs = [
        upload_image,
        l_min_slider,
        l_max_slider,
        tonality_min_slider,
        tonality_max_slider,
        chroma_slider,
        skin_checkbox,
        eye_checkbox,
        hair_checkbox,
    ]
    outputs = [
        processed_image_output,
        l_hist_output,
        tonality_hist_output,
        chroma_hist_output,
        chroma_chart_output,
        undertone_chart_output,
        overtone_chart_output,
        tonality_chart_output,
        season_chart_output,
    ]

    # Set up change event triggers for the sliders
    l_min_slider.change(process_image, inputs=inputs, outputs=outputs)
    l_max_slider.change(process_image, inputs=inputs, outputs=outputs)
    tonality_min_slider.change(process_image, inputs=inputs, outputs=outputs)
    tonality_max_slider.change(process_image, inputs=inputs, outputs=outputs)
    chroma_slider.change(process_image, inputs=inputs, outputs=outputs)

    # upload_image.change(process_image, inputs=inputs, outputs=outputs)
    skin_checkbox.change(process_image, inputs=inputs, outputs=outputs)
    eye_checkbox.change(process_image, inputs=inputs, outputs=outputs)
    hair_checkbox.change(process_image, inputs=inputs, outputs=outputs)

    # Link the submit button to the analyze_image function
    submit_button.click(process_image, inputs=inputs, outputs=outputs)


demo.launch()