update app.py

app.py

@@ -1,103 +1,117 @@
+# app.py
 import gradio as gr
 from ultralytics import YOLO
 import cv2
 import numpy as np
 import torch
-import wikipedia
 from PIL import Image
 import pandas as pd
 import os
 import uuid
 from datetime import datetime
-from h3 import h3
-
-# Initialize models
-model = YOLO("yolov8n.pt")  # You can change this to a fine-tuned model
-midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
+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").small_transform
+midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms", trust_repo=True).small_transform
 
-# CSV path
+# ========================
+# CSV Initialization
+# ========================
 csv_file = "tree_measurements.csv"
 if not os.path.exists(csv_file):
-    pd.DataFrame(columns=[
-        "Timestamp", "Latitude", "Longitude", "H3_Index", 
-        "Estimated_Height", "Species", "Image_File"
-    ]).to_csv(csv_file, index=False)
+    pd.DataFrame(columns=["Timestamp", "Estimated_Height", "Species", "Lat", "Lon", "H3_Index", "Image_File"]).to_csv(csv_file, index=False)
 
-# Convert lat/lon to H3 index
-def latlon_to_h3(lat, lon, resolution=9):
-    return h3.geo_to_h3(lat, lon, resolution)
+# Dummy Tree Classifier (replace with a real model or API later)
+def classify_species(image):
+    return "Unknown Species"
 
-# Main processing function
-def process_tree_image(image, latitude, longitude):
+# ========================
+# Tree Processing Function
+# ========================
+def process_tree(image, lat, lon):
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
-    # Object detection
-    results = model(img_cv)
+    # Run YOLOv8 detection
+    results = yolo_model(img_cv)
     detections = results[0].boxes.data.cpu().numpy()
-
     if len(detections) == 0:
-        return "No tree detected", image
+        return "No tree detected.", None, None
 
-    # Crop the first detected tree
+    # Crop first detected tree
     x1, y1, x2, y2, conf, cls = detections[0]
     tree_crop = img_cv[int(y1):int(y2), int(x1):int(x2)]
 
-    # Depth estimation
+    # 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)
 
-    # Get H3 location index
-    h3_index = latlon_to_h3(latitude, longitude)
+    # Species classification
+    species = classify_species(image)
 
-    # Wikipedia species info
-    try:
-        species_guess = wikipedia.summary("tree", sentences=1)
-    except:
-        species_guess = "Could not retrieve species info"
+    # Location and H3
+    h3_index = h3.geo_to_h3(lat, lon, 9)
 
-    # Save image with unique name
+    # Save image
     image_id = f"tree_{uuid.uuid4().hex[:8]}.png"
     image.save(image_id)
 
-    # Append to CSV
-    new_data = pd.DataFrame([{
+    # Save to CSV
+    new_entry = pd.DataFrame([{
         "Timestamp": timestamp,
-        "Latitude": latitude,
-        "Longitude": longitude,
-        "H3_Index": h3_index,
         "Estimated_Height": approx_height,
-        "Species": species_guess,
+        "Species": species,
+        "Lat": lat,
+        "Lon": lon,
+        "H3_Index": h3_index,
         "Image_File": image_id
     }])
-    new_data.to_csv(csv_file, mode='a', index=False, header=False)
-
-    return (
-        f"📏 Estimated Tree Height: {approx_height} meters\n"
-        f"🌍 H3 Location Index: {h3_index}\n"
-        f"🌿 Species Info: {species_guess}",
-        image
+    new_entry.to_csv(csv_file, mode='a', header=False, index=False)
+
+    return f"Height: {approx_height} meters\nSpecies: {species}", Image.fromarray(tree_crop), generate_map()
+
+# ========================
+# Folium Map from CSV
+# ========================
+def generate_map():
+    df = pd.read_csv(csv_file)
+    fmap = folium.Map(location=[20, 78], zoom_start=5)
+    for _, row in df.iterrows():
+        folium.Marker(
+            location=[row["Lat"], row["Lon"]],
+            popup=f"{row['Timestamp']}\n{row['Species']}\n{row['Estimated_Height']} m"
+        ).add_to(fmap)
+    fmap.save("map.html")
+    return "map.html"
+
+# ========================
+# Gradio UI
+# ========================
+with gr.Blocks() as demo:
+    gr.Markdown("## 🌳 Tree Height & Species Estimator")
+    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]
     )
 
-# Gradio UI
-interface = gr.Interface(
-    fn=process_tree_image,
-    inputs=[
-        gr.Image(type="pil", label="📸 Capture Tree Image"),
-        gr.Number(label="🌐 Latitude"),
-        gr.Number(label="🌐 Longitude")
-    ],
-    outputs=[
-        gr.Textbox(label="📊 Tree Info"),
-        gr.Image(label="🖼️ Captured Image")
-    ],
-    title="🌳 Tree Height & Species Estimator",
-    description="Use your webcam or upload an image of a tree. Enter your GPS location to get height, species info, and H3 geolocation. All data is saved to CSV."
-)
-
-interface.launch()
+# ========================
+demo.launch()