import streamlit as st
from PIL import Image
from utils.model_loader import load_all_models
from utils.predictor import predict_tree, predict_species, predict_stage
from utils.gradcam import (
generate_gradcam_pytorch_resnet,
generate_gradcam_pytorch_mobilenet,
generate_gradcam_keras
)
st.set_page_config(
page_title="TreeVision AI",
page_icon="๐ณ",
layout="wide",
initial_sidebar_state="expanded"
)
st.markdown("""
""", unsafe_allow_html=True)
# โโ Model Loading โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
@st.cache_resource(show_spinner=False)
def get_models():
return load_all_models()
with st.spinner("๐ฟ Initializing TreeVision AI..."):
try:
models = get_models()
except FileNotFoundError as e:
st.error(str(e))
st.stop()
except Exception as e:
st.error(f"Failed to load models: {e}")
st.stop()
# โโ Sidebar โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
with st.sidebar:
st.markdown("## ๐ณ TreeVision AI")
st.markdown("---")
st.markdown("#### How It Works")
st.markdown("""
1
Tree Detection
Is it a tree?
2
Species Detection
Mango or White Gum?
3
Stage Classification
What growth stage?
""", unsafe_allow_html=True)
st.markdown("---")
st.markdown("#### Supported Species")
st.markdown("๐ฅญ **Mango Tree**")
st.markdown("๐ฟ **White Gum (Eucalyptus)**")
st.markdown("---")
st.markdown("#### Growth Stages")
stages = [("๐ฑ", "Seedling"), ("๐ฟ", "Sapling"), ("๐ณ", "Mature"), ("๐", "Overmature")]
for icon, stage in stages:
st.markdown(f"{icon} {stage}")
st.markdown("---")
st.markdown("#### Models Used")
st.caption("Stage 1: ResNet50 (PyTorch) ยท 92.86% acc")
st.caption("Stage 2: MobileNetV2 (TensorFlow) ยท 87.34% acc")
st.caption("Stage 3: MobileNetV2 (PyTorch) ยท 95โ100% acc")
# โโ Hero Section โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
st.markdown("""
๐ฟ AI-Powered ยท Computer Vision
TreeVision AI
Upload a tree photograph for instant species identification and growth stage classification โ powered by a 4-model deep learning pipeline with Grad-CAM visual explanations.
""", unsafe_allow_html=True)
# โโ Upload โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
uploaded_file = st.file_uploader(
"Drop your tree image here",
type=["jpg", "jpeg", "png"],
help="Upload a clear photograph of a single tree for best results."
)
if uploaded_file is None:
col_a, col_b, col_c = st.columns(3)
with col_a:
st.info("๐ธ **Step 1** โ Upload a tree photograph")
with col_b:
st.info("๐ค **Step 2** โ AI pipeline runs automatically")
with col_c:
st.info("๐ **Step 3** โ View results with Grad-CAM heatmaps")
st.stop()
try:
image = Image.open(uploaded_file).convert("RGB")
except Exception as e:
st.error(f"Could not open image: {e}")
st.stop()
st.markdown("
", unsafe_allow_html=True)
col_img, col_info = st.columns([1, 1], gap="large")
with col_img:
st.markdown("YOUR IMAGE
", unsafe_allow_html=True)
st.image(image, use_container_width=True)
with col_info:
st.markdown("#### Pipeline Overview")
st.markdown("""
This image will be analyzed in three sequential stages. Each stage produces a classification result and a Grad-CAM heatmap showing which regions of the image influenced the decision.
**Grad-CAM** (Gradient-weighted Class Activation Mapping) highlights the areas the neural network focused on โ red means high attention, blue means low attention.
""")
st.markdown("""
1
Tree vs Non-Tree โ ResNet50
2
Mango vs White Gum โ MobileNetV2
3
Growth Stage โ Species-specific model
""", unsafe_allow_html=True)
st.markdown("
", unsafe_allow_html=True)
# โโ Stage 1: Tree Detection โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
st.markdown("""
๐ Stage 1
Tree Detection
ResNet50 determines whether the uploaded image contains a tree.
""", unsafe_allow_html=True)
with st.spinner("Analyzing image for tree presence..."):
try:
tree_label, tree_conf, tree_tensor = predict_tree(
image, models["tree_vs_nontree"], models["device"]
)
tree_heatmap = generate_gradcam_pytorch_resnet(
image, models["tree_vs_nontree"], tree_tensor
)
except Exception as e:
st.error(f"Error in tree detection: {e}")
st.stop()
col1, col2 = st.columns(2, gap="large")
with col1:
st.image(image, use_container_width=True)
st.markdown("ORIGINAL IMAGE
", unsafe_allow_html=True)
with col2:
st.image(tree_heatmap, use_container_width=True)
st.markdown("GRAD-CAM HEATMAP
", unsafe_allow_html=True)
if tree_label == "Non-Tree":
st.markdown(f"""
๐ซ No Tree Detected
The model classified this image as Non-Tree with {tree_conf:.1%} confidence.
The pipeline has stopped. Please upload a clear photograph of a tree.
""", unsafe_allow_html=True)
st.stop()
conf_pct = int(tree_conf * 100)
st.markdown(f"""
โ
Detection Result
Tree Detected
Model confidence: {tree_conf:.1%}
""", unsafe_allow_html=True)
st.markdown("
", unsafe_allow_html=True)
# โโ Stage 2: Species Detection โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
st.markdown("""
๐ฟ Stage 2
Species Detection
MobileNetV2 (TensorFlow) identifies whether this is a Mango or White Gum tree.
""", unsafe_allow_html=True)
st.markdown("""
โ ๏ธ Note: This system is trained exclusively on Mango and White Gum (Eucalyptus) trees.
Accuracy may be reduced for other species.
""", unsafe_allow_html=True)
with st.spinner("Identifying tree species..."):
try:
species_label, species_conf, species_array = predict_species(
image, models["species"]
)
species_class_idx = 1 if species_label == "Mango" else 0
species_heatmap = generate_gradcam_keras(
image, models["species"], species_array, species_class_idx
)
except Exception as e:
st.error(f"Error in species detection: {e}")
st.stop()
col3, col4 = st.columns(2, gap="large")
with col3:
st.image(image, use_container_width=True)
st.markdown("ORIGINAL IMAGE
", unsafe_allow_html=True)
with col4:
st.image(species_heatmap, use_container_width=True)
st.markdown("GRAD-CAM HEATMAP
", unsafe_allow_html=True)
species_icon = "๐ฅญ" if species_label == "Mango" else "๐ฟ"
conf_pct2 = int(species_conf * 100)
st.markdown(f"""
๐ฟ Species Result
{species_icon} {species_label} Tree
Model confidence: {species_conf:.1%}
""", unsafe_allow_html=True)
st.markdown("
", unsafe_allow_html=True)
# โโ Stage 3: Growth Stage โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
st.markdown(f"""
๐ฑ Stage 3
Growth Stage Classification
Species-specific MobileNetV2 model classifies the tree's growth stage.
""", unsafe_allow_html=True)
if species_label == "Mango":
stage_model = models["mango_stage"]
model_used = "Mango Stage Model (95.83% accuracy)"
else:
stage_model = models["gum_stage"]
model_used = "White Gum Stage Model (100% accuracy)"
st.caption(f"๐ง Using: {model_used}")
with st.spinner("Classifying growth stage..."):
try:
stage_label, stage_conf, stage_tensor = predict_stage(
image, stage_model, models["device"]
)
stage_heatmap = generate_gradcam_pytorch_mobilenet(
image, stage_model, stage_tensor
)
except Exception as e:
st.error(f"Error in stage classification: {e}")
st.stop()
col5, col6 = st.columns(2, gap="large")
with col5:
st.image(image, use_container_width=True)
st.markdown("ORIGINAL IMAGE
", unsafe_allow_html=True)
with col6:
st.image(stage_heatmap, use_container_width=True)
st.markdown("GRAD-CAM HEATMAP
", unsafe_allow_html=True)
stage_icons = {"Seedling": "๐ฑ", "Sapling": "๐ฟ", "Mature": "๐ณ", "Overmature": "๐"}
stage_icon = stage_icons.get(stage_label, "๐ฟ")
conf_pct3 = int(stage_conf * 100)
st.markdown(f"""
๐ฑ Growth Stage
{stage_icon} {stage_label}
Model confidence: {stage_conf:.1%}
""", unsafe_allow_html=True)
# โโ Summary โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
st.markdown(f"""
๐ Classification Summary
๐ณ
Detection
Tree
{tree_conf:.1%} confidence
{species_icon}
Species
{species_label}
{species_conf:.1%} confidence
{stage_icon}
Growth Stage
{stage_label}
{stage_conf:.1%} confidence
""", unsafe_allow_html=True)