Update app.py

app.py

@@ -1,75 +1,103 @@
 import gradio as gr
 from ultralytics import YOLO
-import torch
 import cv2
 import numpy as np
+import torch
 import wikipedia
 from PIL import Image
-
-# Load YOLO model for tree detection
-yolo_model = YOLO("yolov8n.pt")
-
-# Load MiDaS model for depth estimation
+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")
 midas.to("cpu").eval()
-
-# ✅ FIXED: Use correct MiDaS transform attribute
 midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms").small_transform
 
-def estimate_tree_height(image):
-    # Convert to OpenCV format
-    image_np = np.array(image)
-    image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+# CSV path
+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)
 
-    # Run YOLO detection
-    results = yolo_model(image_bgr)
-    boxes = results[0].boxes.xyxy.cpu().numpy()
-    labels = results[0].boxes.cls.cpu().numpy()
+# Convert lat/lon to H3 index
+def latlon_to_h3(lat, lon, resolution=9):
+    return h3.geo_to_h3(lat, lon, resolution)
 
-    # Filter for trees (assuming class 0 is tree - adjust if needed)
-    tree_boxes = [box for box, label in zip(boxes, labels) if int(label) == 0]
+# Main processing function
+def process_tree_image(image, latitude, longitude):
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
-    if not tree_boxes:
-        return "No tree detected", None, None
+    # Object detection
+    results = model(img_cv)
+    detections = results[0].boxes.data.cpu().numpy()
 
-    x1, y1, x2, y2 = map(int, tree_boxes[0])
-    tree_crop = image_np[y1:y2, x1:x2]
+    if len(detections) == 0:
+        return "No tree detected", image
+
+    # Crop the 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
-    input_tensor = midas_transforms(Image.fromarray(image_np)).to("cpu")
+    input_tensor = midas_transforms(Image.fromarray(cv2.cvtColor(tree_crop, cv2.COLOR_BGR2RGB))).unsqueeze(0)
     with torch.no_grad():
-        depth_map = midas(input_tensor.unsqueeze(0))[0]
-
-    depth_resized = torch.nn.functional.interpolate(
-        depth_map.unsqueeze(0),
-        size=image_np.shape[:2],
-        mode="bicubic",
-        align_corners=False
-    ).squeeze().cpu().numpy()
+        depth = midas(input_tensor).squeeze().cpu().numpy()
+    approx_height = round(np.max(depth) - np.min(depth), 2)
 
-    avg_depth = np.mean(depth_resized[y1:y2, x1:x2])
-    estimated_height_m = avg_depth * 1.8  # scale arbitrarily for demo
+    # Get H3 location index
+    h3_index = latlon_to_h3(latitude, longitude)
 
-    # Wikipedia summary
+    # Wikipedia species info
     try:
-        summary = wikipedia.summary("tree", sentences=2)
-    except Exception:
-        summary = "Tree species information not available."
-
-    return f"Estimated Tree Height: {estimated_height_m:.2f} meters", Image.fromarray(tree_crop), summary
-
-# Gradio Interface
-demo = gr.Interface(
-    fn=estimate_tree_height,
-    inputs=gr.Image(type="pil"),
+        species_guess = wikipedia.summary("tree", sentences=1)
+    except:
+        species_guess = "Could not retrieve species info"
+
+    # Save image with unique name
+    image_id = f"tree_{uuid.uuid4().hex[:8]}.png"
+    image.save(image_id)
+
+    # Append to CSV
+    new_data = pd.DataFrame([{
+        "Timestamp": timestamp,
+        "Latitude": latitude,
+        "Longitude": longitude,
+        "H3_Index": h3_index,
+        "Estimated_Height": approx_height,
+        "Species": species_guess,
+        "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
+    )
+
+# 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 Height Estimate"),
-        gr.Image(label="Detected Tree"),
-        gr.Textbox(label="Tree Species Info")
+        gr.Textbox(label="📊 Tree Info"),
+        gr.Image(label="🖼️ Captured Image")
     ],
-    title="🌳 Tree Measurement App",
-    description="Upload or capture a tree image to estimate its height and get basic species info."
+    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."
 )
-demo.launch()
 
+interface.launch()