app.py

import gradio as gr
import pandas as pd
from PIL import Image
import torch
import torchvision.transforms as T
import os
import json
import sentence_transformers
from huggingface_hub import hf_hub_download
import pickle
import timm
import google.generativeai as genai

# ============================================
# 1. LOAD IMAGE CLASSIFICATION MODEL
# ============================================
print("Loading image classification model...")
REPO_ID = "keerthikoganti/architecture-design-stages-compact-cnn"

pkl_path = hf_hub_download(repo_id=REPO_ID, filename="model_bundle.pkl")
with open(pkl_path, "rb") as f:
    bundle = pickle.load(f)

architecture = bundle["architecture"]
num_classes = bundle["num_classes"]
class_names = bundle["class_names"]
state_dict = bundle["state_dict"]

device = "cuda" if torch.cuda.is_available() else "cpu"
model = timm.create_model(architecture, pretrained=False, num_classes=num_classes)
model.load_state_dict(state_dict)
model.eval().to(device)

TFM = T.Compose([
    T.Resize(224), 
    T.CenterCrop(224), 
    T.ToTensor(), 
    T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])

print("✓ Image classification model loaded successfully!")

# ============================================
# 2. LOAD TEXT CLASSIFICATION MODEL
# ============================================
print("Loading text classification model...")
from autogluon.tabular import TabularPredictor
import shutil

text_repo_id = "kaitongg/my-autogluon-model"
download_dir = "downloaded_predictor"

if os.path.exists(download_dir):
    shutil.rmtree(download_dir)

os.makedirs(download_dir, exist_ok=True)

from huggingface_hub import snapshot_download
downloaded_path = snapshot_download(
    repo_id=text_repo_id,
    repo_type="model",
    local_dir=download_dir,
    local_dir_use_symlinks=False,
)

predictor_path = os.path.join(downloaded_path, "autogluon_predictor")
# Bypass Python version check (model trained on 3.12, running on 3.10)
loaded_predictor_from_hub = TabularPredictor.load(
    predictor_path, 
    require_py_version_match=False,
    require_version_match=False
)

embedding_model = sentence_transformers.SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

print("✓ Text classification model loaded successfully!")

# ============================================
# 3. INITIALIZE GEMINI API
# ============================================
print("Initializing Gemini API...")

# Get API key from environment variable (set in Hugging Face Spaces secrets)
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")

if GEMINI_API_KEY:
    genai.configure(api_key=GEMINI_API_KEY)
    
    # Try models in order of preference (using full model paths)
    model_names_to_try = [
        'models/gemini-2.5-flash',           # Latest stable flash model
        'models/gemini-flash-latest',        # Auto-updates to latest
        'models/gemini-2.0-flash',           # Fallback to 2.0
        'models/gemini-pro-latest',          # Pro version
    ]
    
    gemini_model = None
    for model_name in model_names_to_try:
        try:
            gemini_model = genai.GenerativeModel(model_name)
            # Test with a simple query
            test_response = gemini_model.generate_content("Test")
            print(f"✓ Gemini API initialized successfully with {model_name}!")
            break
        except Exception as e:
            print(f"Failed to load {model_name}: {str(e)[:100]}")
            continue
    
    if gemini_model is None:
        print("⚠️ Warning: Could not initialize any Gemini model")
else:
    gemini_model = None
    print("⚠️ Warning: GEMINI_API_KEY not found in environment variables")

# ============================================
# 4. LLM ATTITUDE MAPPING
# ============================================
llm_attitude_mapping = {
    "brainstorm": "creative and encouraging",
    "design_iteration": "constructive and detailed, focusing on improvements",
    "design_optimization": "critical and focused on efficiency and refinement",
    "final_review": "thorough and critical, evaluating completeness and adherence to requirements",
    "random": "neutral and informative, perhaps suggesting a relevant stage",
}

# ============================================
# 5. TEXT CLASSIFICATION FUNCTION
# ============================================
def perform_text_classification_and_format(text: str) -> tuple:
    text_classification_formatted = "No text provided"
    text_classification_probabilities = {"No High Concept": 0.0, "High Concept": 0.0}
    predicted_text_label = "0"

    if text and loaded_predictor_from_hub is not None and embedding_model is not None:
        try:
            embeddings = embedding_model.encode(
                [text],
                batch_size=1,
                show_progress_bar=False,
                convert_to_numpy=True,
                normalize_embeddings=False,
            )

            n, d = embeddings.shape
            text_df_processed = pd.DataFrame(embeddings, columns=[f"e{i}" for i in range(d)])

            text_proba_df = loaded_predictor_from_hub.predict_proba(text_df_processed)

            text_classification_probabilities = {
                "No High Concept": float(text_proba_df.iloc[0].get("0", 0.0)),
                "High Concept": float(text_proba_df.iloc[0].get("1", 0.0)),
            }

            predicted_text_label = str(loaded_predictor_from_hub.predict(text_df_processed).iloc[0])

            if predicted_text_label == "1":
                has_high_concept = "Yes"
                confidence = text_classification_probabilities["High Concept"]
            else:
                has_high_concept = "No"
                confidence = text_classification_probabilities["No High Concept"]

            text_classification_formatted = f"High Concept: {has_high_concept} (Confidence: {confidence:.2f})"

        except Exception as e:
            print(f"Error processing text: {e}")
            text_classification_formatted = f"Text classification failed: {e}"

    return text_classification_formatted, text_classification_probabilities, predicted_text_label

# ============================================
# 6. COMBINED CLASSIFICATION FUNCTION
# ============================================
def perform_classification_and_format(image: Image.Image, text: str) -> tuple:
    image_classification_results = {"error": "No image provided"}
    design_stage = "unknown"

    if image is not None and model is not None:
        try:
            img_tensor = TFM(image).unsqueeze(0).to(device)

            with torch.no_grad():
                img_output = model(img_tensor)

            img_probabilities = torch.softmax(img_output, dim=1)[0]
            predicted_class_index = torch.argmax(img_probabilities).item()
            design_stage = class_names[predicted_class_index]

            image_classification_results = {
                class_names[i]: float(img_probabilities[i]) 
                for i in range(len(class_names))
            }

            print(f"✓ Image classified as: {design_stage}")

        except Exception as e:
            print(f"❌ Error processing image: {e}")
            image_classification_results = {"error": f"Image classification failed: {e}"}

    text_classification_formatted, text_classification_probabilities, predicted_text_label = perform_text_classification_and_format(text)

    return image_classification_results, text_classification_probabilities, text_classification_formatted, predicted_text_label

# ============================================
# 7. PROMPT GENERATION FUNCTION
# ============================================
def generate_prompt_only(image_classification_results: dict, 
                         text_classification_probabilities: dict, 
                         predicted_text_label: str, 
                         text: str) -> str:
    design_stage = "unknown"
    if image_classification_results and "error" not in image_classification_results:
        try:
            design_stage = max(image_classification_results, key=image_classification_results.get)
        except Exception:
            design_stage = "unknown"

    has_high_concept = "Unable to determine"
    confidence = 0.0
    if text_classification_probabilities and "error" not in text_classification_probabilities:
        try:
            if predicted_text_label == "1":
                has_high_concept = "Yes"
                confidence = text_classification_probabilities.get("High Concept", 0.0)
            else:
                has_high_concept = "No"
                confidence = text_classification_probabilities.get("No High Concept", 0.0)
        except Exception:
            has_high_concept = "Unable to determine"
            confidence = 0.0

    llm_attitude = llm_attitude_mapping.get(design_stage, llm_attitude_mapping["random"])

    prompt = f"""You are an architecture education assistant helping a student understand architectural concepts.

Context:
- Design stage: {design_stage}
- Your feedback style should be: {llm_attitude}
- Abstract concepts detected: {has_high_concept} (confidence: {confidence:.2f})

Student's input: "{text}"

Instructions:
Please provide educational feedback in 250-350 words that:
1. Uses simple, everyday examples and analogies
2. Explains any abstract architectural concepts in accessible language
3. Provides specific, actionable suggestions for improvement
4. Maintains an encouraging yet constructive tone
5. Ends with a complete sentence

Focus on being helpful and educational rather than critical.
"""
    return prompt

# ============================================
# 8. GEMINI FEEDBACK GENERATION
# ============================================
def generate_feedback_from_prompt(prompt_input: str) -> str:
    if gemini_model is None:
        return "⚠️ Gemini API not configured. Please set GEMINI_API_KEY in Hugging Face Spaces secrets."
    
    try:
        print("Generating feedback with Gemini...")
        
        # Extract just the user's input text
        user_text = prompt_input
        if "Student's input:" in prompt_input:
            parts = prompt_input.split("Student's input:")
            if len(parts) > 1:
                user_text = parts[1].strip().strip('"')
        
        # Ultra-simplified prompt - just the core request
        simple_prompt = f"Provide brief educational feedback on this architectural description: {user_text}"
        
        print(f"Sending prompt ({len(simple_prompt)} chars)")
        
        # Minimal configuration - only what's absolutely necessary
        response = gemini_model.generate_content(simple_prompt)
        
        print(f"Response received")
        
        # Extract text
        llm_response_text = None
        
        try:
            llm_response_text = response.text
            print(f"✓ Got text ({len(llm_response_text)} chars)")
        except Exception as e:
            print(f"Failed to get text: {str(e)[:100]}")
            
            # Try alternative extraction
            if response.candidates and response.candidates[0].content:
                candidate = response.candidates[0]
                if candidate.content.parts:
                    texts = [part.text for part in candidate.content.parts if hasattr(part, 'text')]
                    if texts:
                        llm_response_text = "".join(texts)
                        print(f"✓ Got text from parts ({len(llm_response_text)} chars)")
        
        if not llm_response_text:
            return "⚠️ No response generated. This may be an API limitation. Try:\n- Shorter, simpler descriptions\n- Removing technical terms\n- Testing with basic input like 'large windows'"
        
        return llm_response_text.strip()
        
    except Exception as e:
        error_msg = str(e)
        print(f"Error: {error_msg[:200]}")
        return f"Error: {error_msg}"

# ============================================
# 9. GRADIO INTERFACE
# ============================================
examples = [
    ["https://huggingface.co/datasets/kaitongg/image/resolve/main/5e848c2d622e7abe1ad48504_5e01ce9f0d272014d0353cd1_Things-You-Need-to-Organize-a-3D-Rendering-Architectural-Project-EASY-RENDER.jpeg", 
     "Exploring spatial relationships and material palettes."],
    ["https://huggingface.co/datasets/kaitongg/image/resolve/main/EXISTING-FIRST-FLOOR-PRES-scaled-e1635965923983.jpg", 
     "The window size is too small."],
    ["https://huggingface.co/datasets/kaitongg/image/resolve/main/bilbao_sketch.png", 
     "The facade expresses the building's relationship with the urban context."],
]

with gr.Blocks(css="""
    .left-column, .middle-column, .right-column {min-width: 300px !important;}
    .textbox-container textarea {min-height: 150px !important;}
""") as demo:
    
    gr.Markdown("# 🏛️ Architecture Feedback Generator (Powered by Gemini)")
    gr.Markdown("""
    Upload an architectural image and provide a text description or question. 
    The system will classify the design stage, analyze the text for high-level concepts, 
    generate a customized prompt, and provide AI-powered feedback using Google's Gemini.
    """)

    with gr.Row():
        # LEFT COLUMN - Input Section
        with gr.Column(scale=1, elem_classes="left-column"):
            gr.Markdown("### 📥 Input")
            image_input = gr.Image(type="pil", label="Upload Architectural Image", height=300)
            text_input = gr.Textbox(
                label="Enter Text Description or Question", 
                placeholder="Describe your architectural design, ask questions, or provide context...",
                lines=6,
                elem_classes="textbox-container"
            )
            classify_button = gr.Button("🔍 Classify & Generate Prompt", variant="primary", size="lg")
        
        # MIDDLE COLUMN - Classification & Prompt Section
        with gr.Column(scale=1, elem_classes="middle-column"):
            gr.Markdown("### 📊 Classification Results & Prompt")
            image_output_label = gr.Label(
                num_top_classes=len(class_names), 
                label="Image Classification (Design Stage)"
            )
            text_output_textbox = gr.Textbox(
                label="Text Classification (High Concept Detection)", 
                lines=2,
                elem_classes="textbox-container"
            )
            prompt_output_textbox = gr.Textbox(
                label="Generated Prompt (Editable)", 
                lines=10,
                interactive=True,
                elem_classes="textbox-container"
            )
            generate_feedback_button = gr.Button("✨ Generate AI Feedback", variant="primary", size="lg")
        
        # RIGHT COLUMN - Gemini Output Section
        with gr.Column(scale=1, elem_classes="right-column"):
            gr.Markdown("### 🤖 AI-Generated Feedback")
            llm_output_text = gr.Textbox(
                label="Gemini Response", 
                lines=20,
                elem_classes="textbox-container",
                show_copy_button=True
            )

    # Hidden state variables
    text_classification_probabilities_state = gr.State()
    predicted_text_label_state = gr.State()

    # Step 1: Classification
    classification_outputs = classify_button.click(
        fn=perform_classification_and_format,
        inputs=[image_input, text_input],
        outputs=[
            image_output_label, 
            text_classification_probabilities_state, 
            text_output_textbox,
            predicted_text_label_state
        ]
    )

    # Step 2: Generate Prompt
    def generate_prompt_wrapper(img_res, txt_prob, predicted_label, txt):
        return generate_prompt_only(img_res, txt_prob, predicted_label, txt)

    classification_outputs.then(
        fn=generate_prompt_wrapper,
        inputs=[
            image_output_label, 
            text_classification_probabilities_state, 
            predicted_text_label_state,
            text_input
        ],
        outputs=prompt_output_textbox
    )

    # Step 3: Gemini Feedback
    generate_feedback_button.click(
        fn=generate_feedback_from_prompt,
        inputs=[prompt_output_textbox],
        outputs=llm_output_text
    )

    # Examples Section
    gr.Markdown("---")
    gr.Markdown("### 💡 Example Inputs")
    
    def generate_full_chain_output(img, txt):
        img_res, txt_prob, txt_fmt, predicted_label = perform_classification_and_format(img, txt)
        prompt = generate_prompt_only(img_res, txt_prob, predicted_label, txt)
        llm_res = generate_feedback_from_prompt(prompt)
        return img_res, txt_fmt, prompt, llm_res

    gr.Examples(
        examples=examples,
        inputs=[image_input, text_input],
        outputs=[image_output_label, text_output_textbox, prompt_output_textbox, llm_output_text],
        fn=generate_full_chain_output,
        cache_examples=False
    )

if __name__ == "__main__":
    demo.launch()