import streamlit as st

# Custom CSS for better styling
st.markdown("""
    <style>
        .main-title {
            font-size: 36px;
            color: #4A90E2;
            font-weight: bold;
            text-align: center;
        }
        .sub-title {
            font-size: 24px;
            color: #4A90E2;
            margin-top: 20px;
        }
        .section {
            background-color: #f9f9f9;
            padding: 15px;
            border-radius: 10px;
            margin-top: 20px;
        }
        .section h2 {
            font-size: 22px;
            color: #4A90E2;
        }
        .section p, .section ul {
            color: #666666;
        }
        .link {
            color: #4A90E2;
            text-decoration: none;
        }
        .benchmark-table {
            width: 100%;
            border-collapse: collapse;
            margin-top: 20px;
        }
        .benchmark-table th, .benchmark-table td {
            border: 1px solid #ddd;
            padding: 8px;
            text-align: left;
        }
        .benchmark-table th {
            background-color: #4A90E2;
            color: white;
        }
        .benchmark-table td {
            background-color: #f2f2f2;
        }
    </style>
""", unsafe_allow_html=True)

# Main Title
st.markdown('<div class="main-title">ConvNeXT Image Classification</div>', unsafe_allow_html=True)

# Description
st.markdown("""
<div class="section">
    <p><strong>ConvNeXT</strong> is a state-of-the-art image classification model developed by Facebook. The model <strong>ConvNextForImageClassification</strong> can load ConvNeXT models that compete favorably with Transformers in terms of accuracy and scalability, achieving 87.8% ImageNet top-1 accuracy and outperforming Swin Transformers on COCO detection and ADE20K segmentation, while maintaining the simplicity and efficiency of standard ConvNets.</p>
    <p>This annotator is compatible with all the models trained/fine-tuned by using ConvNextForImageClassification for PyTorch or TFConvNextForImageClassification for TensorFlow models in Hugging Face.</p>
    <p>The model used in this demo is <code>image_classifier_convnext_tiny_224_local</code>, adapted from Hugging Face and curated for scalability and production-readiness using Spark NLP.</p>
</div>
""", unsafe_allow_html=True)

# Image Classification Overview
st.markdown('<div class="sub-title">What is Image Classification?</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <p><strong>Image Classification</strong> is a computer vision task where an algorithm is trained to recognize and classify objects within images. This process involves assigning a label or category to an image based on its visual content.</p>
    <h2>How It Works</h2>
    <p>Image classification typically involves the following steps:</p>
    <ul>
        <li><strong>Data Collection</strong>: Gather a dataset of labeled images.</li>
        <li><strong>Preprocessing</strong>: Normalize and resize images to prepare them for the model.</li>
        <li><strong>Model Training</strong>: Use a machine learning model, such as ConvNeXT, to learn patterns and features from the images.</li>
        <li><strong>Inference</strong>: Apply the trained model to new images to predict their labels.</li>
    </ul>
    <h2>Why Use Image Classification?</h2>
    <p>Image classification can automate and streamline many tasks, such as:</p>
    <ul>
        <li>Identifying objects in photos for content tagging.</li>
        <li>Enhancing search functionality by categorizing images.</li>
        <li>Supporting autonomous systems like self-driving cars.</li>
    </ul>
    <h2>Applications</h2>
    <p>Applications of image classification span across various industries:</p>
    <ul>
        <li><strong>Healthcare</strong>: Diagnosing diseases from medical images.</li>
        <li><strong>Retail</strong>: Sorting and tagging product images.</li>
        <li><strong>Security</strong>: Facial recognition for authentication.</li>
    </ul>
    <h2>Importance</h2>
    <p>Image classification is crucial because it enables machines to interpret visual data, which is essential for creating intelligent systems capable of understanding and interacting with the world in a more human-like manner.</p>
    <p>The <strong>ConvNeXT</strong> model used in this example is a state-of-the-art approach for image classification, offering advanced performance and scalability. It utilizes convolutional architecture to capture intricate patterns and relationships within images, enhancing classification accuracy and efficiency.</p>
</div>
""", unsafe_allow_html=True)

# How to Use
st.markdown('<div class="sub-title">How to Use the Model</div>', unsafe_allow_html=True)
st.code('''
import sparknlp
from sparknlp.base import *
from sparknlp.annotator import *
from pyspark.ml import Pipeline

# Load image data
imageDF = spark.read \\
    .format("image") \\
    .option("dropInvalid", value = True) \\
    .load("src/test/resources/image/")

# Define Image Assembler
imageAssembler = ImageAssembler() \\
    .setInputCol("image") \\
    .setOutputCol("image_assembler")

# Define ConvNeXT classifier
imageClassifier = ConvNextForImageClassification \\
    .pretrained("image_classifier_convnext_tiny_224_local", "en") \\
    .setInputCols(["image_assembler"]) \\
    .setOutputCol("class")

# Create pipeline
pipeline = Pipeline().setStages([imageAssembler, imageClassifier])

# Apply pipeline to image data
pipelineDF = pipeline.fit(imageDF).transform(imageDF)

# Show results
pipelineDF \\
  .selectExpr("reverse(split(image.origin, '/'))[0] as image_name", "class.result") \\
  .show(truncate=False)
''', language='python')

# Results
st.markdown('<div class="sub-title">Results</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <table class="benchmark-table">
        <tr>
            <th>Image Name</th>
            <th>Result</th>
        </tr>
        <tr>
            <td>dog.JPEG</td>
            <td>[whippet]</td>
        </tr>
        <tr>
            <td>cat.JPEG</td>
            <td>[Siamese]</td>
        </tr>
        <tr>
            <td>bird.JPEG</td>
            <td>[peacock]</td>
        </tr>
    </table>
</div>
""", unsafe_allow_html=True)

# Model Information
st.markdown('<div class="sub-title">Model Information</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <table class="benchmark-table">
        <tr>
            <th>Attribute</th>
            <th>Description</th>
        </tr>
        <tr>
            <td><strong>Model Name</strong></td>
            <td>image_classifier_convnext_tiny_224_local</td>
        </tr>
        <tr>
            <td><strong>Compatibility</strong></td>
            <td>Spark NLP 5.0.0+</td>
        </tr>
        <tr>
            <td><strong>License</strong></td>
            <td>Open Source</td>
        </tr>
        <tr>
            <td><strong>Edition</strong></td>
            <td>Official</td>
        </tr>
        <tr>
            <td><strong>Input Labels</strong></td>
            <td>[image_assembler]</td>
        </tr>
        <tr>
            <td><strong>Output Labels</strong></td>
            <td>[class]</td>
        </tr>
        <tr>
            <td><strong>Language</strong></td>
            <td>en</td>
        </tr>
        <tr>
            <td><strong>Size</strong></td>
            <td>107.6 MB</td>
        </tr>
    </table>
</div>
""", unsafe_allow_html=True)

# Predicted Entities
st.markdown('<div class="sub-title">Predicted Entities</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <ul>
        <li>turnstile</li>
        <li>damselfly</li>
        <li>mixing bowl</li>
        <li>sea snake</li>
        <li>cockroach</li>
        <li>...and many more</li>
    </ul>
</div>
""", unsafe_allow_html=True)

# Data Source Section
st.markdown('<div class="sub-title">Data Source</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <p>The ConvNeXT model is available on <a class="link" href="https://huggingface.co/models" target="_blank">Hugging Face</a>. This model was trained on a large dataset of images and can be used for accurate image classification.</p>
</div>
""", unsafe_allow_html=True)

# References
st.markdown('<div class="sub-title">References</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <ul>
        <li><a class="link" href="https://sparknlp.org/2023/07/05/image_classifier_convnext_tiny_224_local_en.html" target="_blank" rel="noopener">ConvNeXT Model on Spark NLP</a></li>
        <li><a class="link" href="https://huggingface.co/facebook/convnext-tiny-224" target="_blank" rel="noopener">ConvNeXT Model on Hugging Face</a></li>
        <li><a class="link" href="https://github.com/facebookresearch/ConvNeXT" target="_blank" rel="noopener">ConvNeXT GitHub Repository</a></li>
        <li><a class="link" href="https://arxiv.org/abs/2201.03545" target="_blank" rel="noopener">ConvNeXT Paper</a></li>
    </ul>
</div>
""", unsafe_allow_html=True)

# Community & Support
st.markdown('<div class="sub-title">Community & Support</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <ul>
        <li><a class="link" href="https://sparknlp.org/" target="_blank">Official Website</a>: Documentation and examples</li>
        <li><a class="link" href="https://join.slack.com/t/spark-nlp/shared_invite/zt-198dipu77-L3UWNe_AJ8xqDk0ivmih5Q" target="_blank">Slack</a>: Live discussion with the community and team</li>
        <li><a class="link" href="https://github.com/JohnSnowLabs/spark-nlp" target="_blank">GitHub</a>: Bug reports, feature requests, and contributions</li>
        <li><a class="link" href="https://medium.com/spark-nlp" target="_blank">Medium</a>: Spark NLP articles</li>
        <li><a class="link" href="https://www.youtube.com/channel/UCmFOjlpYEhxf_wJUDuz6xxQ/videos" target="_blank">YouTube</a>: Video tutorials</li>
    </ul>
</div>
""", unsafe_allow_html=True)