update app.py

app.py

@@ -1,117 +1,129 @@
-# app.py
 import gradio as gr
 from ultralytics import YOLO
+import torch
 import cv2
 import numpy as np
-import torch
 from PIL import Image
 import pandas as pd
 import os
 import uuid
 from datetime import datetime
-import h3
 import folium
-
-# ========================
-# Load Models
-# ========================
-yolo_model = YOLO("yolov8n.pt")
-midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", trust_repo=True)
-midas.to("cpu").eval()
-midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms", trust_repo=True).small_transform
-
-# ========================
-# CSV Initialization
-# ========================
+from h3 import h3
+import base64
+
+# Load YOLO model for tree detection
+model = YOLO("yolov8n.pt")
+
+# Try loading MiDaS depth model
+try:
+    midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", trust_repo=True)
+    midas.to("cpu").eval()
+    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms", trust_repo=True).small_transform
+    use_depth = True
+except Exception as e:
+    print(f"Depth model load failed: {e}")
+    use_depth = False
+
+# CSV file setup
 csv_file = "tree_measurements.csv"
 if not os.path.exists(csv_file):
-    pd.DataFrame(columns=["Timestamp", "Estimated_Height", "Species", "Lat", "Lon", "H3_Index", "Image_File"]).to_csv(csv_file, index=False)
+    pd.DataFrame(columns=["Timestamp", "Latitude", "Longitude", "H3_Index", "Height", "Species", "Image_File"]).to_csv(csv_file, index=False)
 
-# Dummy Tree Classifier (replace with a real model or API later)
-def classify_species(image):
-    return "Unknown Species"
+# Dummy classifier
+def classify_tree_species(image):
+    # Placeholder - returns fixed label
+    return "Generic Tree"
 
-# ========================
-# Tree Processing Function
-# ========================
-def process_tree(image, lat, lon):
+# Process function
+def analyze_tree(image, latitude, longitude):
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    image_np = np.array(image)
+    img_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
 
-    # Run YOLOv8 detection
-    results = yolo_model(img_cv)
+    results = model(img_cv)
     detections = results[0].boxes.data.cpu().numpy()
+
     if len(detections) == 0:
-        return "No tree detected.", None, None
+        return "No tree detected", image, "N/A", generate_map()
 
-    # Crop first detected tree
+    # First detected tree
     x1, y1, x2, y2, conf, cls = detections[0]
-    tree_crop = img_cv[int(y1):int(y2), int(x1):int(x2)]
-
-    # Estimate height from depth
-    input_tensor = midas_transforms(Image.fromarray(cv2.cvtColor(tree_crop, cv2.COLOR_BGR2RGB))).unsqueeze(0)
-    with torch.no_grad():
-        depth = midas(input_tensor).squeeze().cpu().numpy()
-    approx_height = round(np.max(depth) - np.min(depth), 2)
-
-    # Species classification
-    species = classify_species(image)
-
-    # Location and H3
-    h3_index = h3.geo_to_h3(lat, lon, 9)
+    crop = img_cv[int(y1):int(y2), int(x1):int(x2)]
+    tree_crop = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
+
+    # Depth estimation
+    if use_depth:
+        try:
+            input_tensor = midas_transforms(tree_crop).unsqueeze(0)
+            with torch.no_grad():
+                depth = midas(input_tensor).squeeze().cpu().numpy()
+            approx_height = round(np.max(depth) - np.min(depth), 2)
+        except Exception:
+            approx_height = "Unavailable"
+    else:
+        approx_height = "Unavailable"
+
+    # Geolocation + H3 index
+    h3_index = h3.geo_to_h3(float(latitude), float(longitude), 9)
+
+    # Species
+    species = classify_tree_species(tree_crop)
 
     # Save image
     image_id = f"tree_{uuid.uuid4().hex[:8]}.png"
-    image.save(image_id)
+    tree_crop.save(image_id)
 
-    # Save to CSV
-    new_entry = pd.DataFrame([{
+    # Append to CSV
+    pd.DataFrame([{
         "Timestamp": timestamp,
-        "Estimated_Height": approx_height,
-        "Species": species,
-        "Lat": lat,
-        "Lon": lon,
+        "Latitude": latitude,
+        "Longitude": longitude,
         "H3_Index": h3_index,
+        "Height": approx_height,
+        "Species": species,
         "Image_File": image_id
-    }])
-    new_entry.to_csv(csv_file, mode='a', header=False, index=False)
+    }]).to_csv(csv_file, mode='a', header=False, index=False)
 
-    return f"Height: {approx_height} meters\nSpecies: {species}", Image.fromarray(tree_crop), generate_map()
+    return f"Height: {approx_height} units\nSpecies: {species}", tree_crop, species, generate_map()
 
-# ========================
-# Folium Map from CSV
-# ========================
+# Render map
 def generate_map():
+    if not os.path.exists(csv_file):
+        return "No map yet."
+
     df = pd.read_csv(csv_file)
-    fmap = folium.Map(location=[20, 78], zoom_start=5)
+    if df.empty:
+        return "No map data."
+
+    lat, lon = df.iloc[-1][["Latitude", "Longitude"]]
+    fmap = folium.Map(location=[lat, lon], zoom_start=14)
     for _, row in df.iterrows():
         folium.Marker(
-            location=[row["Lat"], row["Lon"]],
-            popup=f"{row['Timestamp']}\n{row['Species']}\n{row['Estimated_Height']} m"
+            location=[row["Latitude"], row["Longitude"]],
+            popup=f"{row['Species']} ({row['Height']} units)"
         ).add_to(fmap)
     fmap.save("map.html")
-    return "map.html"
+    with open("map.html", "r", encoding="utf-8") as f:
+        return f.read()
 
-# ========================
 # Gradio UI
-# ========================
 with gr.Blocks() as demo:
-    gr.Markdown("## 🌳 Tree Height & Species Estimator")
+    gr.Markdown("## 🌳 Tree Height & Species Estimator with Map & Logger")
+
     with gr.Row():
-        image_input = gr.Image(type="pil", label="Capture/Upload Tree Photo")
-        lat_input = gr.Number(label="Latitude")
-        lon_input = gr.Number(label="Longitude")
-    submit_btn = gr.Button("Estimate Height & Species")
-    output_text = gr.Textbox(label="Result")
-    output_image = gr.Image(label="Detected Tree")
-    output_map = gr.HTML(label="Tree Map")
-
-    submit_btn.click(
-        fn=process_tree,
-        inputs=[image_input, lat_input, lon_input],
-        outputs=[output_text, output_image, output_map]
-    )
-
-# ========================
-demo.launch()
+        image_input = gr.Image(type="pil", label="📸 Tree Image")
+        lat_input = gr.Textbox(label="🌍 Latitude", placeholder="e.g., 12.9716")
+        lon_input = gr.Textbox(label="🌍 Longitude", placeholder="e.g., 77.5946")
 
+    btn = gr.Button("Analyze Tree")
+
+    output_text = gr.Textbox(label="📏 Results")
+    output_crop = gr.Image(label="🌲 Detected Tree")
+    output_species = gr.Textbox(label="🌳 Species")
+    map_html = gr.HTML(label="🗺️ Tree Map")
+
+    btn.click(analyze_tree, inputs=[image_input, lat_input, lon_input],
+              outputs=[output_text, output_crop, output_species, map_html])
+
+demo.launch()