# Coded by Manith Jayaba , All Rights Reserved ©
import streamlit as st
from PIL import Image
import torch
from torchvision import transforms
import cv2
import numpy as np
from facenet_pytorch import MTCNN

# Function to load the ViT model
def load_model(model_path):
    model = torch.load(model_path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    return model, device

# Initialize MTCNN for face detection
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
mtcnn = MTCNN(keep_all=True, device=device)

# Function to preprocess the image using MTCNN for face detection
def preprocess_image(image, device):
    # Convert PIL image to OpenCV format
    open_cv_image = np.array(image.convert("RGB"))
    # Convert RGB to BGR for OpenCV
    open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)

    # Convert OpenCV image back to PIL Image for MTCNN
    pil_image = Image.fromarray(cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2RGB))

    # Use MTCNN to detect faces
    boxes, _ = mtcnn.detect(pil_image)
    if boxes is not None:
        # Crop the first detected face (for simplicity)
        box = boxes[0].astype(int)
        cropped_face = open_cv_image[box[1]:box[3], box[0]:box[2]]
        # Convert cropped face back to PIL for further processing
        processed_image = Image.fromarray(cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB))
    else:
        st.warning("No face Detected/Low Quality Image , Results can be Inaccurate")
        processed_image = image  # Use the original image if no face is detected

    # Transform image for model
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    image_tensor = transform(processed_image).to(device)
    image_tensor = image_tensor.unsqueeze(0)  # Add a batch dimension
    return image_tensor, processed_image

# Function for inference
def predict(image_tensor, model, device):
    model.eval()
    with torch.no_grad():
        outputs = model(image_tensor)
        # Check for 'logits' or 'predictions' attribute and use accordingly
        logits = outputs.logits if hasattr(outputs, 'logits') else outputs.predictions
        probabilities = torch.nn.functional.softmax(logits, dim=1)
        predicted_class = torch.argmax(probabilities, dim=1)
    return predicted_class, probabilities

# Streamlit UI setup
st.title("SkinDet: Oily/Dry Skin Level Predictor 👨🏻‍🔬🧬")
st.markdown("### Utilized with Cutting-Edge `Vision Transformers`")
st.markdown("#### Designed & Developed by `Manith Jayaba` & `Shakya Dissanayake`")
st.image('skin1.jpg', caption='Skin ANN',use_column_width=True)

st.write("""
This app can measure the oiliness and dryness of your skin in to 5 levels:
1. Level 1: Normal Oily Skin
2. Level 2: Mid Oily Skin
3. Level 3: High Oily Skin
4. Level 4: Very High Oily Skin
5. High Level Dry Skin

`Disclaimer:`
- This application provides insights on skin oiliness for informational purposes only, not medical advice.
- Consult a dermatologist for personalized skincare assessment and recommendations.
- Accuracy varies with image quality and individual skin characteristics.
- Classifications rely on machine learning and may not always match the specific skin condition.
- Users are responsible for interpreting and acting on application insights, exercising caution in skincare decisions.
- Professional medical advice is recommended for comprehensive evaluation and treatment.
""")

# Coded by Manith Jayaba , All Rights Reserved ©

model_path = "model_v2.0.pt"  # Adjust this path as necessary
model, device = load_model(model_path)

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    #st.image(image, caption='Uploaded Image', use_column_width=True)

    # Preprocess the image and perform inference
    image_tensor, final_image = preprocess_image(image, device)
    predicted_class, probabilities = predict(image_tensor, model, device)
    # {'Dry': 0, 'Normal': 1, 'Oily': 2}

    nprobabilities = probabilities.numpy()
    #st.write(f"Predicted class: {predicted_class.item()}")

    dry_percentage = "{:.2f}".format(nprobabilities[0][0] * 100)
    oily_percentage = "{:.2f}".format(nprobabilities[0][2] * 100)
    normal_percentage = "{:.2f}".format(nprobabilities[0][1] * 100)

    # Display progress bars with correctly formatted percentage values
    st.write(f"Dry Skin: {dry_percentage}%")
    st.progress(int(float(dry_percentage)))
    st.write(f"Oily Skin: {oily_percentage}%")
    st.progress(int(float(oily_percentage)))
    st.write(f"Normal Skin: {normal_percentage}%")
    st.progress(int(float(normal_percentage)))

    # if predicted_class.item() == 2:
    #     st.write(f"Oily Skin : {nprobabilities[0][2]}")
    # elif predicted_class.item() == 1:
    #     st.write(f"Normal Skin : {nprobabilities[0][1]}")
    # elif predicted_class.item() == 0:
    #     st.write(f"Dry Skin : {nprobabilities[0][0]}")
        
    #st.write(f"Probabilities: {nprobabilities}")

    if float(oily_percentage) >= 80:
            st.warning('Level 4: Very High Oiliness Level')
            st.warning("To manage very high oiliness in your skin, start with a gentle, oil-free cleanser twice daily, followed by an alcohol-free toner. Opt for a lightweight, non-comedogenic moisturizer to hydrate without adding grease, and use a broad-spectrum sunscreen daily. Incorporate weekly exfoliation with salicylic acid to keep pores clear and occasionally use clay masks to absorb excess oil. Blotting papers can help manage shine throughout the day. Remember, a balanced diet and proper hydration also play a crucial role in controlling oil production. Adjust your skincare routine based on how your skin responds, and consider consulting a dermatologist for personalized advice or treatments.")
    elif 60 <= float(oily_percentage) < 80:
        st.warning('Level 3: High Oiliness Level')
        st.warning("To manage high oiliness, adopt a straightforward skincare routine: Start with a gentle, oil-free cleanser twice a day, followed by an alcohol-free toner to balance skin pH. Use a lightweight, non-comedogenic moisturizer to hydrate without adding shine. Apply a broad-spectrum, oil-free sunscreen daily. Weekly, gently exfoliate to prevent clogged pores and use a clay mask to absorb excess oil. Keep blotting papers on hand for midday touch-ups. Adjust your diet to minimize high-sugar and fatty foods, which can influence oil production. Regularly assess your skin's reaction to products and adjust as needed, considering a dermatologist's advice for persistent issues.")
    elif 50 < float(oily_percentage) < 60:
        st.info('Level 2: Mid Oiliness Level')
        st.info("For managing mid-level oiliness, balance is key. Start with a mild cleanser to wash away excess oil without stripping the skin, doing so morning and night. Incorporate a lightweight, non-comedogenic moisturizer to hydrate and protect the skin without adding greasiness. A broad-spectrum, oil-free sunscreen is essential for daily protection. Use a gentle exfoliant once or twice a week to remove dead skin cells and keep pores clear, opting for products with salicylic acid for its oil-dissolving properties. Consider using a clay mask weekly to draw out impurities and control shine. Adjust your routine based on seasonal changes or hormonal fluctuations, as these can affect your skin’s oil production. A balanced diet, rich in fruits, vegetables, and plenty of water, supports overall skin health. Monitor how your skin responds to different products and tweak your routine as necessary, seeking a dermatologist’s advice for stubborn or severe issues.")
    elif float(oily_percentage) <= 50 and float(dry_percentage)>= 55:
        st.warning('High Level Dry Skin')
        st.warning("For high-level dry skin, a nurturing and moisture-rich routine is essential. Begin with a hydrating, soap-free cleanser to gently remove impurities without stripping skin of its natural oils. Follow up with a rich, emollient moisturizer that contains ingredients like hyaluronic acid, glycerin, or ceramides to lock in moisture and strengthen the skin's barrier. Applying an oil-based serum before your moisturizer can offer an additional hydration boost. Don’t forget to apply a broad-spectrum sunscreen during the day to protect from UV rays, which can exacerbate dryness. Once or twice a week, use a gentle exfoliant to remove dead skin cells, improving absorption of moisturizers. Consider incorporating a weekly hydrating mask for deep moisture replenishment. To support skin hydration from the inside, drink plenty of water and consider using a humidifier in dry environments. Regularly assess your skin’s condition and adjust your skincare products as needed, especially with seasonal changes. For persistent dryness or irritation, consult a dermatologist for tailored advice.")
    elif float(oily_percentage) <= 50:
        st.info('Level 1: Normal Oiliness Level')
        st.info("For normal oiliness levels, a simple and balanced skincare routine works best. Begin with a gentle cleanser to remove dirt and impurities without over-drying, using it both in the morning and evening. Follow up with a lightweight moisturizer that maintains hydration without contributing to excess oil. Daily sunscreen application is crucial to protect your skin from UV damage; opt for a formula that feels comfortable on your skin type. Incorporate a mild exfoliant once a week to slough off dead skin cells and promote cell turnover, ensuring your skin remains clear and vibrant. Adjust your skincare products seasonally, as your skin's needs may change with the weather. Alongside your skincare routine, maintaining a healthy diet and staying hydrated contribute to the overall appearance and health of your skin. Regularly assess how your skin responds to your routine, and don't hesitate to tweak it as needed for the best results.")
    

    # Display the final processed image
    st.image(final_image, caption='Processed Image', use_column_width=True)

    # Coded by Manith Jayaba , All Rights Reserved ©