app.py

"""
Netra AI - Construction Material Classifier
MVP Demo using Claude Vision API
"""

import gradio as gr
import anthropic
import os
from PIL import Image
import base64
from io import BytesIO

# Material classification prompt
CLASSIFICATION_PROMPT = """
You are a construction material classifier for Netra AI. Classify this image into ONE of these 4 classes:

1. **Reet (Sand)**: All grades of sand - from fine powdery sand to coarse gritty sand. Shows smooth to granular texture.

2. **12mm VSI**: Uniform, small, cubical "clean" stones. Very consistent size (~10-15mm). Manufactured aggregate with sharp edges.

3. **Stone**: Large, irregular boulders. Significant size variation, deep voids between rocks. Raw, unprocessed appearance.

4. **GSB (Graded Stone Base)**: Mixed-size material with rocks, gravel, grit AND fine dust/sand filling gaps. Key indicator is fine dust matrix.

Respond in this exact format:
CLASS: [class name]
CONFIDENCE: [High/Medium/Low]
REASONING: [One sentence explanation of visual features that led to this classification]
"""


def classify_material(image):
    """Classify construction material from image using Netra AI"""

    if image is None:
        return "Please upload an image", "", ""

    # Get API key from environment
    api_key = os.getenv("ANTHROPIC_API_KEY", "")
    if not api_key:
        return "❌ System Error", "", "AI service not configured. Please contact support."

    try:
        # Configure Claude
        client = anthropic.Anthropic(api_key=api_key)

        # Convert to PIL Image if needed
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)

        # Convert image to base64
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        image_data = base64.b64encode(buffered.getvalue()).decode("utf-8")

        # Generate classification
        response = client.messages.create(
            model="claude-opus-4-6",
            max_tokens=150,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": "image/jpeg",
                                "data": image_data,
                            },
                        },
                        {"type": "text", "text": CLASSIFICATION_PROMPT}
                    ],
                }
            ],
        )
        result_text = response.content[0].text.strip()

        # Parse response
        lines = result_text.split('\n')
        class_name = ""
        confidence = ""
        reasoning = ""

        for line in lines:
            if line.startswith("CLASS:"):
                class_name = line.replace("CLASS:", "").strip()
            elif line.startswith("CONFIDENCE:"):
                confidence = line.replace("CONFIDENCE:", "").strip()
            elif line.startswith("REASONING:"):
                reasoning = line.replace("REASONING:", "").strip()

        # Format output
        if class_name:
            result = f"✅ RESULT:\n\n{class_name}"
        else:
            result = "⏳ AWAITING CLASSIFICATION"

        confidence_display = f"Confidence: {confidence}" if confidence else ""
        reasoning_display = f"💡 **Analysis:** {reasoning}" if reasoning else result_text

        return result, confidence_display, reasoning_display

    except Exception as e:
        error_msg = str(e)
        if "API_KEY_INVALID" in error_msg or "401" in error_msg:
            return "❌ Invalid API Key", "", "Please check your Gemini API key"
        elif "quota" in error_msg.lower():
            return "❌ API Quota Exceeded", "", "Your API key has exceeded its quota"
        else:
            return f"❌ Error: {error_msg[:100]}", "", "Please try again"


# Custom CSS
custom_css = """
#header {
    text-align: center;
    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
    padding: 20px;
    border-radius: 10px;
    color: white;
}
#result {
    font-size: 72px;
    font-weight: 900;
    text-align: center;
    padding: 50px 30px;
    border-radius: 20px;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    margin-bottom: 25px;
    box-shadow: 0 8px 16px rgba(0,0,0,0.2);
    text-transform: uppercase;
    letter-spacing: 2px;
    line-height: 1.2;
}
#confidence {
    font-size: 24px;
    font-weight: 600;
    text-align: center;
    padding: 15px;
    background: #f0f9ff;
    border-radius: 10px;
    margin-bottom: 20px;
}
#reasoning {
    font-size: 14px;
    text-align: left;
    padding: 15px;
    background: #f9fafb;
    border-radius: 8px;
    margin-bottom: 20px;
    color: #4b5563;
}
#about {
    font-size: 13px;
    color: #6b7280;
    padding-top: 20px;
    border-top: 1px solid #e5e7eb;
}
"""

# Build Gradio interface
with gr.Blocks(css=custom_css, title="Netra AI - Material Classifier") as demo:

    gr.Markdown(
        """
        <div id="header">
        <h1>🏗️ Netra AI - Construction Material Classifier</h1>
        <p>AI-powered material identification for construction sites</p>
        </div>
        """,
        elem_id="header"
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📸 Upload Material Image")
            image_input = gr.Image(
                type="pil",
                label="Construction Material Photo",
                height=400
            )

            classify_btn = gr.Button("🚀 Classify Material", variant="primary", size="lg")

            gr.Markdown(
                """
                **Supported Materials:**
                - 🏖️ Reet (Sand)
                - 🔷 12mm VSI (Uniform Aggregate)
                - 🪨 Stone (Large Boulders)
                - 🏗️ GSB (Mixed Graded Base)
                """
            )

            gr.Markdown("### 📷 Try Sample Images")
            gr.Examples(
                examples=[
                    ["sanity_test_dataset/12mm_VSI_01.jpeg"],
                    ["sanity_test_dataset/Reet_01_fine.jpeg"],
                    ["sanity_test_dataset/Stone_01.jpeg"],
                    ["sanity_test_dataset/GSB_01.jpeg"],
                ],
                inputs=image_input,
                label=""
            )

        with gr.Column(scale=1):
            result_output = gr.Markdown("", elem_id="result")
            confidence_output = gr.Markdown("", elem_id="confidence")
            reasoning_output = gr.Markdown("", elem_id="reasoning")

            gr.Markdown(
                """
                <div id="about">

                ### 📊 About This Demo

                This is a proof-of-concept for **Netra AI's** automated material classification system.

                **Technology:**
                - Vision AI for real-time material identification
                - Prevents grade fraud at construction sites
                - Automates dispatch logging

                **Use Cases:**
                - Gate monitoring at crusher plants
                - Dispatch verification
                - Quality control

                ---
                **Netra AI** - *Making every tonne count.*

                For commercial inquiries: [Contact Us](mailto:guptaraghu321@gmail.com)

                </div>
                """
            )


    # Connect button
    classify_btn.click(
        fn=classify_material,
        inputs=[image_input],
        outputs=[result_output, confidence_output, reasoning_output]
    )

# Launch
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )