import streamlit as st
import cv2
import numpy as np
from PIL import Image, ExifTags
from mtcnn import MTCNN
from deepface import DeepFace
from transformers import BlipProcessor, BlipForConditionalGeneration

# Initialize MTCNN for face detection
detector = MTCNN()

# Load BLIP model for scene description
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# Function to auto-rotate the image based on EXIF data
def auto_rotate_image(image):
    try:
        for orientation in ExifTags.TAGS.keys():
            if ExifTags.TAGS[orientation] == 'Orientation':
                break

        exif = image._getexif()
        if exif is not None:
            orientation_value = exif.get(orientation)
            if orientation_value == 3:
                image = image.rotate(180, expand=True)
            elif orientation_value == 6:
                image = image.rotate(270, expand=True)
            elif orientation_value == 8:
                image = image.rotate(90, expand=True)
    except (AttributeError, KeyError, IndexError, TypeError):
        pass  # No EXIF data or issues reading EXIF, skip rotation

    return image

# Function to generate a scene description using BLIP
def describe_scene(image):
    inputs = processor(image, return_tensors="pt")
    outputs = blip_model.generate(**inputs)
    caption = processor.decode(outputs[0], skip_special_tokens=True)
    return caption

# Function for face analysis using DeepFace
def analyze_face(face_image):
    face_array = np.array(face_image)
    try:
        result = DeepFace.analyze(face_array, actions=['gender', 'age', 'race', 'emotion'], enforce_detection=False)
        return result
    except Exception as e:
        return {"error": str(e)}

# Main Streamlit application
def main():
    st.title("Comprehensive Image Analysis")
    st.write("Upload an image to get a scene description and detailed face analysis (gender, age, race, emotions).")

    # Upload image
    uploaded_image = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png", "webp"])

    if uploaded_image is not None:
        image = Image.open(uploaded_image).convert("RGB")

        # Auto-rotate image if possible
        image = auto_rotate_image(image)

        st.image(image, caption="Uploaded Image", use_container_width=True)

        # Option for manual rotation
        rotation_angle = st.slider("Rotate Image (Degrees)", -180, 180, 0)
        if rotation_angle != 0:
            image = image.rotate(rotation_angle, expand=True)
            st.image(image, caption="Rotated Image", use_container_width=True)

        if st.button("Analyze"):
            with st.spinner("Analyzing..."):
                # Scene description
                scene_caption = describe_scene(image)
                st.subheader("Scene Description")
                st.write(scene_caption)

                # Detect faces using MTCNN
                faces = detector.detect_faces(np.array(image))

                if faces:
                    # Limit to only the first detected face
                    st.write("**One Face Detected. Analyzing...**")

                    # Get the first face
                    face = faces[0]
                    x, y, width, height = face['box']
                    face_crop = image.crop((x, y, x + width, y + height))
                    face_crop_resized = face_crop.resize((150, 150))

                    st.subheader("Face Analysis")
                    st.image(face_crop_resized, caption="Detected Face", width=150)

                    # Analyze face with DeepFace
                    result = analyze_face(face_crop)
                    if "error" in result:
                        st.error(f"Face Analysis Error: {result['error']}")
                    else:
                        single_result = result[0] if isinstance(result, list) else result

                        # Gender
                        dominant_gender = single_result.get('dominant_gender', 'N/A')
                        gender_confidence = int(round(single_result.get('gender', {}).get(dominant_gender, 0)))
                        st.write(f"**Gender:** {dominant_gender.capitalize()} ({gender_confidence}%)")

                        # Age
                        age = single_result.get('age', 'N/A')
                        st.write(f"**Age:** {age}")

                        # Race
                        race = single_result.get('race', {})
                        race_text = ", ".join([f"{r}: {int(round(p))}%" for r, p in race.items()])
                        st.write(f"**Race:** {race_text}")

                        # Emotions
                        emotions = single_result.get('emotion', {})
                        emotion_text = ", ".join([f"{e}: {int(round(p))}%" for e, p in emotions.items()])
                        st.write(f"**Emotions:** {emotion_text}")
                else:
                    st.write("No faces detected.")

if __name__ == "__main__":
    main()