Update app.py

app.py

@@ -1,128 +1,102 @@
+# arboreal_gradio_app/app.py
 import gradio as gr
-from ultralytics import YOLO
-import torch
 import cv2
+import torch
 import numpy as np
+import time
+import wikipedia
 import pandas as pd
-import h3
-import folium
-from folium.plugins import MarkerCluster
-from PIL import Image
-import base64
 import os
+from ultralytics import YOLO
+from PIL import Image
 from datetime import datetime
 
-# Load YOLO model
-model = YOLO("yolov8n.pt")
-
-# Try loading MiDaS depth estimation model
-try:
-    midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", trust_repo=True)
-    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms", trust_repo=True)
-    transform = midas_transforms.small_transform
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    midas.to(device)
-    height_estimation_available = True
-except:
-    height_estimation_available = False
-
-# Placeholder species classifier (could be replaced by real model)
-def classify_species(image):
-    return "Neem Tree"  # Dummy result
-
-# Create output CSV if it doesn't exist
-if not os.path.exists("tree_data.csv"):
-    pd.DataFrame(columns=["timestamp", "latitude", "longitude", "h3_index", "species", "height_m", "image_path"]).to_csv("tree_data.csv", index=False)
-
-# Tree analysis function
-def analyze_tree(image, latitude, longitude):
-    image_np = np.array(image)
-
-    # Detect tree using YOLO
-    results = model(image_np)[0]
-    boxes = results.boxes.xyxy.cpu().numpy().astype(int)
-    classes = results.boxes.cls.cpu().numpy().astype(int)
-
-    height_m = "N/A"
-
-    # Depth estimation (if available)
-    if height_estimation_available:
-        input_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        input_image = cv2.resize(input_image, (256, 256))
-        input_tensor = transform(input_image).to(device)
-        with torch.no_grad():
-            prediction = midas(input_tensor.unsqueeze(0))
-            prediction = torch.nn.functional.interpolate(
-                prediction.unsqueeze(1),
-                size=image_np.shape[:2],
-                mode="bicubic",
-                align_corners=False,
-            ).squeeze()
-            depth_map = prediction.cpu().numpy()
-            height_m = round(float(np.median(depth_map)), 2)
-
-    # Species classification (placeholder)
-    species = classify_species(image)
-
-    # H3 geolocation
-    h3_index = h3.geo_to_h3(float(latitude), float(longitude), 9)
-
-    # Save image
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    image_path = f"tree_{timestamp}.jpg"
-    image.save(image_path)
-
-    # Append to CSV
-    df = pd.read_csv("tree_data.csv")
-    df.loc[len(df)] = [timestamp, latitude, longitude, h3_index, species, height_m, image_path]
-    df.to_csv("tree_data.csv", index=False)
-
-    return f"Species: {species}\nEstimated Height: {height_m} meters\nH3 Index: {h3_index}"
-
-# Generate map with markers
-def generate_map():
-    df = pd.read_csv("tree_data.csv")
-    m = folium.Map(location=[20.5937, 78.9629], zoom_start=5)
-    marker_cluster = MarkerCluster().add_to(m)
-
-    for _, row in df.iterrows():
-        folium.Marker(
-            location=[row["latitude"], row["longitude"]],
-            popup=f"Species: {row['species']}\nHeight: {row['height_m']}\nH3: {row['h3_index']}",
-        ).add_to(marker_cluster)
-
-    map_path = "tree_map.html"
-    m.save(map_path)
-    with open(map_path, "r") as f:
-        html_content = f.read()
-    return html_content
-
-# Gradio interface
-def js_location():
-    return """
-    async () => {
-        const pos = await new Promise((resolve, reject) => {
-            navigator.geolocation.getCurrentPosition(resolve, reject);
-        });
-        return [pos.coords.latitude, pos.coords.longitude];
+# Load YOLOv8 model
+model = YOLO("yolov8n.pt")  # You can replace with a custom-trained tree model
+
+# Constants
+KNOWN_PERSON_HEIGHT = 1.7  # meters
+CSV_FILE = "tree_log.csv"
+
+# Initialize CSV if it doesn't exist
+if not os.path.exists(CSV_FILE):
+    pd.DataFrame(columns=["Image", "Height (m)", "Species", "Latitude", "Longitude", "Timestamp"]).to_csv(CSV_FILE, index=False)
+
+# Function to estimate height using YOLO detections
+def estimate_tree_height(img):
+    results = model(img)
+    person_box, tree_box = None, None
+
+    for result in results:
+        for box in result.boxes:
+            cls = result.names[int(box.cls[0])]
+            xyxy = box.xyxy[0].cpu().numpy()
+            if cls.lower() == "person" and person_box is None:
+                person_box = xyxy
+            elif cls.lower() == "tree" and tree_box is None:
+                tree_box = xyxy
+
+    if tree_box is None or person_box is None:
+        return None, "Could not detect both tree and person."
+
+    # Estimate height using pixel ratio
+    tree_h = tree_box[3] - tree_box[1]
+    person_h = person_box[3] - person_box[1]
+    ratio = tree_h / person_h
+    estimated_height = ratio * KNOWN_PERSON_HEIGHT
+
+    return estimated_height, None
+
+# Function to get species info from Wikipedia
+def get_species_info(species_name):
+    try:
+        summary = wikipedia.summary(species_name, sentences=2)
+        return summary
+    except:
+        return "No species info found."
+
+# Main process function
+def process(image, species_guess, latitude, longitude):
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    image_array = np.array(image)
+
+    height, error = estimate_tree_height(image_array)
+    if error:
+        return error, None
+
+    species_info = get_species_info(species_guess)
+
+    # Save to CSV
+    df = pd.read_csv(CSV_FILE)
+    image_filename = f"image_{int(time.time())}.jpg"
+    new_row = {
+        "Image": image_filename,
+        "Height (m)": round(height, 2),
+        "Species": species_guess,
+        "Latitude": latitude,
+        "Longitude": longitude,
+        "Timestamp": timestamp
     }
-    """
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            cam = gr.Image(source="webcam", streaming=True)
-            loc_btn = gr.Button("Get Location")
-            lat = gr.Textbox(label="Latitude")
-            lon = gr.Textbox(label="Longitude")
-            analyze_btn = gr.Button("Analyze Tree")
-            output = gr.Textbox(label="Results")
-        with gr.Column():
-            map_output = gr.HTML()
-
-    loc_btn.click(None, js=js_location(), outputs=[lat, lon])
-    analyze_btn.click(fn=analyze_tree, inputs=[cam, lat, lon], outputs=output)
-    analyze_btn.click(fn=generate_map, outputs=map_output)
-
-if __name__ == "__main__":
-    demo.launch()
+    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
+    df.to_csv(CSV_FILE, index=False)
+
+    return f"Estimated Height: {round(height, 2)} meters\nSpecies: {species_guess}", species_info
+
+# Gradio UI
+iface = gr.Interface(
+    fn=process,
+    inputs=[
+        gr.Image(label="Capture Tree + Person Image"),
+        gr.Textbox(label="Species Guess (optional)", placeholder="E.g. Neem Tree"),
+        gr.Textbox(label="Latitude (optional)", placeholder="Enter from GPS or manually"),
+        gr.Textbox(label="Longitude (optional)", placeholder="Enter from GPS or manually")
+    ],
+    outputs=[
+        gr.Text(label="Result"),
+        gr.Text(label="Species Information (Wikipedia)")
+    ],
+    title="🌲 Tree Analyzer - Height + Species + Location",
+    description="Upload a photo with a tree and a person. Optionally enter location. Returns estimated height and species info. Logs all data to CSV."
+)
+
+iface.launch()