# Ci-Dave from BSCS-AI
# Description: This Python script creates a Streamlit web application for image analysis using computer vision techniques and AI-generated explanations. 
# The app allows users to upload an image, apply edge detection, segmentation, feature extraction, and AI classification. 
# The explanations for each technique are generated using the Gemini API for AI-generated content.

import streamlit as st  # Streamlit library to create the web interface
import numpy as np  # Library for numerical operations
import google.generativeai as genai  # Gemini API for AI-generated explanations

# Random Forest and Logistic Regression model for classification
from sklearn.ensemble import RandomForestClassifier   
from sklearn.linear_model import LogisticRegression   

from skimage.filters import sobel  # Sobel edge detection filter from skimage
from skimage.segmentation import watershed  # Watershed segmentation method
from skimage.feature import canny, hog  # Canny edge detection and HOG feature extraction
from skimage.color import rgb2gray  # Convert RGB images to grayscale

from skimage import io  # I/O functions for reading images
from sklearn.preprocessing import StandardScaler  # Standardization of image data

# Load Gemini API key from Streamlit Secrets configuration
api_key = st.secrets["gemini"]["api_key"]  # Get API key from Streamlit secrets
genai.configure(api_key=api_key)  # Configure the Gemini API with the API keyimport streamlit as st
import numpy as np
import google.generativeai as genai
import matplotlib.pyplot as plt

from sklearn.ensemble import RandomForestClassifier   
from sklearn.linear_model import LogisticRegression   
from skimage.filters import sobel  
from skimage.segmentation import watershed  
from skimage.feature import canny, hog  
from skimage.color import rgb2gray  
from skimage import io  
from sklearn.preprocessing import StandardScaler  
from sklearn.metrics import accuracy_score

# Load Gemini API key
api_key = st.secrets["gemini"]["api_key"]
genai.configure(api_key=api_key)
MODEL_ID = "gemini-1.5-flash"
gen_model = genai.GenerativeModel(MODEL_ID)

def explain_ai(prompt):
    try:
        response = gen_model.generate_content(prompt)
        return response.text
    except Exception as e:
        return f"Error: {str(e)}"


# Sidebar navigation
st.sidebar.title("Navigation")
page = st.sidebar.radio("Go to", ["Home", "Edge Detection", "Segmentation", "Feature Extraction", "AI Classification"])

# Home Page
if page == "Home":
    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
    if uploaded_file is not None:
        image = io.imread(uploaded_file)
        if image.shape[-1] == 4:
            image = image[:, :, :3]
        gray = rgb2gray(image)
        st.image(image, caption="Uploaded Image", use_container_width=True)
        st.session_state["gray"] = gray  # Store for use in other pages

# Edge Detection Page
elif page == "Edge Detection":
    st.title("Edge Detection")
    gray = st.session_state.get("gray")
    if gray is not None:
        edge_method = st.selectbox("Select Edge Detection Method", ["Canny", "Sobel"])
        edges = canny(gray) if edge_method == "Canny" else sobel(gray)
        st.image(edges, caption=f"{edge_method} Edge Detection", use_container_width=True)
        st.text_area("Explanation", explain_ai(f"Explain how {edge_method} edge detection works in computer vision."), height=300)
    else:
        st.warning("Please upload an image on the Home page.")

# Segmentation Page
elif page == "Segmentation":
    st.title("Image Segmentation")
    gray = st.session_state.get("gray")
    if gray is not None:
        seg_method = st.selectbox("Select Segmentation Method", ["Watershed", "Thresholding"])
        if seg_method == "Watershed":
            elevation_map = sobel(gray)
            markers = np.zeros_like(gray)
            markers[gray < 0.3] = 1
            markers[gray > 0.7] = 2
            segmented = watershed(elevation_map, markers.astype(np.int32))
        else:
            threshold_value = st.slider("Choose threshold value", 0, 255, 127)
            segmented = (gray > (threshold_value / 255)).astype(np.uint8) * 255
        st.image(segmented, caption=f"{seg_method} Segmentation", use_container_width=True)
        st.text_area("Explanation", explain_ai(f"Explain how {seg_method} segmentation works in image processing."), height=300)
    else:
        st.warning("Please upload an image on the Home page.")

# Feature Extraction Page
elif page == "Feature Extraction":
    st.title("HOG Feature Extraction")
    gray = st.session_state.get("gray")
    if gray is not None:
        fd, hog_image = hog(gray, pixels_per_cell=(8, 8), cells_per_block=(2, 2), visualize=True)
        st.image(hog_image, caption="HOG Features", use_container_width=True)
        st.text_area("Explanation", explain_ai("Explain how Histogram of Oriented Gradients (HOG) feature extraction works."), height=300)
    else:
        st.warning("Please upload an image on the Home page.")

# AI Classification Page
elif page == "AI Classification":
    st.title("AI Classification")
    gray = st.session_state.get("gray")
    if gray is not None:
        model_choice = st.selectbox("Select AI Model", ["Random Forest", "Logistic Regression"])
        flat_image = gray.flatten().reshape(-1, 1)
        labels = (flat_image > 0.5).astype(int).flatten()
        ai_model = RandomForestClassifier(n_jobs=1) if model_choice == "Random Forest" else LogisticRegression()
        scaler = StandardScaler()
        flat_image_scaled = scaler.fit_transform(flat_image)
        ai_model.fit(flat_image_scaled, labels)
        predictions = ai_model.predict(flat_image_scaled).reshape(gray.shape)
        predictions = (predictions * 255).astype(np.uint8)
        accuracy = accuracy_score(labels, ai_model.predict(flat_image_scaled))
        st.image(predictions, caption=f"{model_choice} Pixel Classification", use_container_width=True)
        st.text_area("Explanation", explain_ai(f"Explain how {model_choice} is used for image classification."), height=300)
        st.write(f"### Accuracy: {accuracy:.2f}")
        fig, ax = plt.subplots()
        ax.bar(["Accuracy"], [accuracy], color='blue')
        ax.set_ylim([0, 1])
        st.pyplot(fig)
    else:
        st.warning("Please upload an image on the Home page.")


MODEL_ID = "gemini-1.5-flash"  # Specify the model ID for Gemini
gen_model = genai.GenerativeModel(MODEL_ID)  # Initialize the Gemini model

# Function to generate explanations using the Gemini API
def explain_ai(prompt):
    """Generate an explanation using Gemini API with error handling."""
    try:
        response = gen_model.generate_content(prompt)  # Get AI-generated content based on prompt
        return response.text  # Return the explanation text
    except Exception as e:
        return f"Error: {str(e)}"  # Return error message if there's an issue

# App title
st.title("Imaize: Smart Image Analyzer with XAI")

# # Image upload section
# uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])  # Allow user to upload an image file

# App Description 
st.markdown("""
This app combines AI-powered image analysis techniques with an easy-to-use interface for explanation generation. 
It leverages advanced computer vision algorithms such as **edge detection**, **image segmentation**, and **feature extraction**. 
Additionally, the app provides **explanations** for each method used, powered by the Gemini API, to make the process more understandable.

The main functionalities of the app include:
- **Edge Detection**: Choose between the Canny and Sobel edge detection methods.
- **Segmentation**: Apply Watershed or Thresholding methods to segment images.
- **Feature Extraction**: Extract Histogram of Oriented Gradients (HOG) features from images.
- **AI Classification**: Classify images using Random Forest or Logistic Regression models.

Whether you're exploring computer vision or simply curious about how these techniques work, this app will guide you through the process with easy-to-understand explanations.
""")

# Instructions on how to use the app 
st.markdown("""
### How to Use the App:

1. **Upload an Image**: Click on the "Upload an image" button to upload an image (in JPG, PNG, or JPEG format) for analysis.
2. **Select Edge Detection**: Choose between **Canny** or **Sobel** edge detection methods. The app will process the image and display the result.
3. **Apply Segmentation**: Select **Watershed** or **Thresholding** segmentation. You can also adjust the threshold for thresholding segmentation.
4. **Extract HOG Features**: Visualize the HOG (Histogram of Oriented Gradients) features from the image.
5. **Choose AI Model for Classification**: Select either **Random Forest** or **Logistic Regression** to classify the image based on pixel information.
6. **Read the Explanations**: For each technique, you'll find a detailed explanation of how it works, powered by AI. Simply read the generated explanation to understand the underlying processes.

### Enjoy exploring and understanding image analysis techniques with AI!
""")

# If an image is uploaded, proceed with the analysis
if uploaded_file is not None:
    image = io.imread(uploaded_file)  # Read the uploaded image using skimage
    if image.shape[-1] == 4:  # If the image has 4 channels (RGBA), remove the alpha channel
        image = image[:, :, :3]

    gray = rgb2gray(image)  # Convert the image to grayscale for processing
    
    st.image(image, caption="Uploaded Image", use_container_width=True)  # Display the uploaded image

    # Edge Detection Section
    st.subheader("Edge Detection")  # Title for edge detection section
    edge_method = st.selectbox("Select Edge Detection Method", ["Canny", "Sobel"], key="edge")  # Select edge detection method
    edges = canny(gray) if edge_method == "Canny" else sobel(gray)  # Apply chosen edge detection method
    edges = (edges * 255).astype(np.uint8)  # Convert edge map to 8-bit image format
    
    col1, col2 = st.columns([1, 1])  # Create two columns for layout
    with col1:
        st.image(edges, caption=f"{edge_method} Edge Detection", use_container_width=True)  # Display the edge detection result
    with col2:
        explanation = explain_ai(f"Explain how {edge_method} edge detection works in computer vision.")  # Get explanation from AI
        st.text_area("Explanation", explanation, height=300)  # Display explanation in a text area

    # Segmentation Section
    st.subheader("Segmentation")  # Title for segmentation section
    seg_method = st.selectbox("Select Segmentation Method", ["Watershed", "Thresholding"], key="seg")  # Select segmentation method
    
    # Perform segmentation based on chosen method
    if seg_method == "Watershed":
        elevation_map = sobel(gray)  # Create elevation map using Sobel filter
        markers = np.zeros_like(gray)  # Initialize marker array
        markers[gray < 0.3] = 1  # Mark low-intensity regions
        markers[gray > 0.7] = 2  # Mark high-intensity regions
        segmented = watershed(elevation_map, markers.astype(np.int32))  # Apply watershed segmentation
    else:
        threshold_value = st.slider("Choose threshold value", 0, 255, 127)  # Slider to choose threshold value
        segmented = (gray > (threshold_value / 255)).astype(np.uint8) * 255  # Apply thresholding segmentation
    
    col1, col2 = st.columns([1, 1])  # Create two columns for layout
    with col1:
        st.image(segmented, caption=f"{seg_method} Segmentation", use_container_width=True)  # Display segmentation result
    with col2:
        explanation = explain_ai(f"Explain how {seg_method} segmentation works in image processing.")  # Get explanation from AI
        st.text_area("Explanation", explanation, height=300)  # Display explanation in a text area

    # HOG Feature Extraction Section
    st.subheader("HOG Feature Extraction")  # Title for HOG feature extraction section
    fd, hog_image = hog(gray, pixels_per_cell=(8, 8), cells_per_block=(2, 2), visualize=True)  # Extract HOG features
    
    col1, col2 = st.columns([1, 1])  # Create two columns for layout
    with col1:
        st.image(hog_image, caption="HOG Features", use_container_width=True)  # Display HOG feature image
    with col2:
        explanation = explain_ai("Explain how Histogram of Oriented Gradients (HOG) feature extraction works.")  # Get explanation from AI
        st.text_area("Explanation", explanation, height=300)  # Display explanation in a text area

    # AI Classification Section
    st.subheader("AI Classification")  # Title for AI classification section
    model_choice = st.selectbox("Select AI Model", ["Random Forest", "Logistic Regression"], key="model")  # Select AI model for classification
    
    flat_image = gray.flatten().reshape(-1, 1)  # Flatten the grayscale image into a 1D array for classification
    labels = (flat_image > 0.5).astype(int).flatten()  # Generate binary labels based on intensity threshold
    
    # Choose model (Random Forest or Logistic Regression)
    ai_model = RandomForestClassifier(n_jobs=1) if model_choice == "Random Forest" else LogisticRegression()  # Initialize the model
    scaler = StandardScaler()  # Standardize the image data for better classification
    flat_image_scaled = scaler.fit_transform(flat_image)  # Scale the image data
    
    ai_model.fit(flat_image_scaled, labels)  # Train the AI model on the image data
    predictions = ai_model.predict(flat_image_scaled).reshape(gray.shape)  # Make predictions on the image
    predictions = (predictions * 255).astype(np.uint8)  # Convert predictions to 8-bit image format
    
    col1, col2 = st.columns([1, 1])  # Create two columns for layout
    with col1:
        st.image(predictions, caption=f"{model_choice} Pixel Classification", use_container_width=True)  # Display classification result
    with col2:
        explanation = explain_ai(f"Explain how {model_choice} is used for image classification.")  # Get explanation from AI
        st.text_area("Explanation", explanation, height=300)  # Display explanation in a text area