Update app.py

app.py

@@ -2,296 +2,275 @@ import gradio as gr
 import torch
 import cv2
 import numpy as np
-import requests
-import json
 from PIL import Image
-import transformers
 from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
 import wikipedia
 import folium
-from geopy.geocoders import Nominatim
-import base64
-from io import BytesIO
 import tempfile
 import os
+import logging
+import warnings
+warnings.filterwarnings("ignore")
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class TreeAnalyzer:
     def __init__(self):
         self.setup_models()
-        self.geolocator = Nominatim(user_agent="tree_analyzer")
         
     def setup_models(self):
-        """Initialize all required models"""
-        print("Loading models...")
+        """Initialize models optimized for HF Spaces"""
+        logger.info("Loading models for HF Spaces...")
         
-        # Load MiDaS model for depth estimation
+        # Load depth estimation model
+        self.midas = None
         try:
-            self.midas = torch.hub.load('intel-isl/MiDaS', 'MiDaS_small')
+            self.midas = torch.hub.load('intel-isl/MiDaS', 'MiDaS_small', trust_repo=True)
             self.midas.eval()
-            self.midas_transforms = torch.hub.load('intel-isl/MiDaS', 'transforms')
+            self.midas_transforms = torch.hub.load('intel-isl/MiDaS', 'transforms', trust_repo=True)
             self.transform = self.midas_transforms.small_transform
-            print("✓ MiDaS model loaded successfully")
+            logger.info("✓ MiDaS loaded")
         except Exception as e:
-            print(f"Error loading MiDaS: {e}")
-            self.midas = None
+            logger.error(f"MiDaS failed: {e}")
             
         # Load plant classification model
-        try:
-            self.plant_classifier = pipeline(
-                "image-classification",
-                model="microsoft/resnet-50",
-                return_top_k=3
-            )
-            print("✓ Plant classifier loaded successfully")
-        except Exception as e:
-            print(f"Error loading plant classifier: {e}")
-            # Fallback to a more specific plant model if available
+        self.plant_classifier = None
+        models_to_try = [
+            "google/vit-base-patch16-224",
+            "microsoft/resnet-50",
+            "facebook/convnext-tiny-224"
+        ]
+        
+        for model_name in models_to_try:
             try:
                 self.plant_classifier = pipeline(
                     "image-classification",
-                    model="google/vit-base-patch16-224",
-                    return_top_k=3
+                    model=model_name,
+                    return_top_k=10
                 )
-                print("✓ Fallback classifier loaded successfully")
-            except:
-                self.plant_classifier = None
-                print("✗ Could not load plant classifier")
+                logger.info(f"✓ Loaded classifier: {model_name}")
+                break
+            except Exception as e:
+                logger.warning(f"Failed to load {model_name}: {e}")
+                continue
     
-    def estimate_tree_height(self, image, known_object_height=1.7):
-        """
-        Estimate tree height using MiDaS depth estimation
-        known_object_height: assumed height of reference object (person = 1.7m)
-        """
+    def estimate_tree_height(self, image):
+        """Estimate tree height using depth estimation"""
         if self.midas is None:
-            return "MiDaS model not available", None
+            return "Height estimation not available (MiDaS model failed to load)"
             
         try:
-            # Convert PIL to OpenCV format
+            # Convert and resize image
             img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+            h, w = img_cv.shape[:2]
+            
+            # Resize for memory efficiency
+            if h > 384 or w > 384:
+                scale = min(384/h, 384/w)
+                new_h, new_w = int(h*scale), int(w*scale)
+                img_cv = cv2.resize(img_cv, (new_w, new_h))
             
-            # Prepare image for MiDaS
-            input_batch = self.transform(img_cv).to(torch.float32)
+            # Process with MiDaS
+            input_batch = self.transform(img_cv)
             
-            # Generate depth map
             with torch.no_grad():
                 prediction = self.midas(input_batch)
                 prediction = torch.nn.functional.interpolate(
                     prediction.unsqueeze(1),
-                    size=img_cv.shape[:2],
+                    size=(img_cv.shape[0], img_cv.shape[1]),
                     mode="bicubic",
                     align_corners=False,
                 ).squeeze()
             
-            # Convert to numpy
             depth_map = prediction.cpu().numpy()
             
-            # Normalize depth map for visualization
-            depth_map_normalized = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
-            depth_map_colored = cv2.applyColorMap(depth_map_normalized, cv2.COLORMAP_PLASMA)
-            
-            # Simple height estimation (this is a simplified approach)
-            # In reality, you'd need camera calibration and more sophisticated methods
-            height, width = depth_map.shape
-            
-            # Assume tree is in the center-upper portion of the image
-            tree_region = depth_map[int(height*0.1):int(height*0.8), int(width*0.3):int(width*0.7)]
-            
-            # Calculate relative height based on depth variations
-            depth_range = np.max(tree_region) - np.min(tree_region)
-            
-            # Rough estimation: scale based on depth range and image dimensions
-            estimated_height = (depth_range / np.max(depth_map)) * height * 0.02  # Scaling factor
-            estimated_height = max(2.0, min(50.0, estimated_height))  # Clamp between 2-50 meters
-            
-            return f"Estimated height: {estimated_height:.1f} meters", depth_map_colored
+            # Simple height estimation
+            h_img, w_img = depth_map.shape
+            center_region = depth_map[h_img//4:3*h_img//4, w_img//4:3*w_img//4]
             
+            if center_region.size > 0:
+                depth_range = np.max(center_region) - np.min(center_region)
+                height_ratio = center_region.shape[0] / h_img
+                estimated_height = max(1.5, min(50.0, (depth_range * height_ratio * 30)))
+                
+                return f"Estimated height: {estimated_height:.1f} meters\n(Approximate estimate based on image depth analysis)"
+            else:
+                return "Could not estimate height from this image"
+                
         except Exception as e:
-            return f"Error in height estimation: {str(e)}", None
+            logger.error(f"Height estimation error: {e}")
+            return f"Height estimation failed: {str(e)}"
     
     def identify_tree_species(self, image):
-        """Identify tree species using image classification"""
+        """Identify tree species with better filtering"""
         if self.plant_classifier is None:
-            return "Plant classifier not available", []
+            return "Species identification not available (classifier failed to load)", []
             
         try:
+            # Resize image for processing
+            if image.size[0] > 224 or image.size[1] > 224:
+                image = image.resize((224, 224), Image.Resampling.LANCZOS)
+            
             # Get predictions
             predictions = self.plant_classifier(image)
             
-            # Filter for tree-related predictions
-            tree_keywords = ['tree', 'oak', 'pine', 'maple', 'birch', 'cedar', 'fir', 'palm', 'willow', 'cherry', 'apple']
-            tree_predictions = []
+            # Enhanced plant/tree keywords
+            plant_keywords = [
+                # Trees
+                'tree', 'oak', 'pine', 'maple', 'birch', 'cedar', 'fir', 'palm', 'willow',
+                'cherry', 'apple', 'spruce', 'poplar', 'ash', 'elm', 'beech', 'sycamore',
+                'acacia', 'eucalyptus', 'magnolia', 'chestnut', 'walnut', 'hickory',
+                'cypress', 'juniper', 'redwood', 'bamboo', 'mahogany', 'teak',
+                # Plants and botanical terms
+                'plant', 'leaf', 'leaves', 'branch', 'bark', 'forest', 'wood', 'botanical',
+                'flora', 'foliage', 'evergreen', 'deciduous', 'conifer', 'hardwood',
+                'softwood', 'timber', 'shrub', 'bush', 'vine', 'fern', 'moss',
+                # Specific species indicators
+                'quercus', 'pinus', 'acer', 'betula', 'fagus', 'tilia', 'fraxinus',
+                'platanus', 'castanea', 'juglans', 'carya', 'ulmus', 'salix'
+            ]
             
+            # Process and score predictions
+            species_candidates = []
             for pred in predictions:
                 label = pred['label'].lower()
-                if any(keyword in label for keyword in tree_keywords):
-                    tree_predictions.append(pred)
-            
-            if not tree_predictions:
-                tree_predictions = predictions[:2]  # Take top 2 if no tree-specific matches
+                confidence = pred['score']
+                
+                # Calculate plant relevance score
+                plant_score = sum(1 for keyword in plant_keywords if keyword in label)
+                is_plant_related = plant_score > 0
+                
+                # Get Wikipedia info
+                wiki_info = self.get_wikipedia_info(pred['label'])
+                
+                species_candidates.append({
+                    'species': pred['label'],
+                    'confidence': confidence,
+                    'plant_score': plant_score,
+                    'is_plant_related': is_plant_related,
+                    'wiki_info': wiki_info
+                })
             
-            # Get Wikipedia information for top prediction
-            species_info = []
-            for pred in tree_predictions[:2]:
-                try:
-                    # Search Wikipedia
-                    wiki_results = wikipedia.search(pred['label'] + " tree", results=1)
-                    if wiki_results:
-                        page = wikipedia.page(wiki_results[0])
-                        summary = wikipedia.summary(wiki_results[0], sentences=2)
-                        species_info.append({
-                            'species': pred['label'],
-                            'confidence': pred['score'],
-                            'wiki_title': page.title,
-                            'summary': summary,
-                            'url': page.url
-                        })
-                except:
-                    species_info.append({
-                        'species': pred['label'],
-                        'confidence': pred['score'],
-                        'wiki_title': 'Not found',
-                        'summary': 'No Wikipedia information available',
-                        'url': None
-                    })
+            # Sort by plant relevance and confidence
+            species_candidates.sort(key=lambda x: (x['plant_score'], x['confidence']), reverse=True)
             
-            return "Species identified successfully", species_info
+            # Return top candidates
+            final_results = species_candidates[:3]
             
+            if any(result['is_plant_related'] for result in final_results):
+                return "Species identification completed", final_results
+            else:
+                return "Possible species identified (may not be plants)", final_results
+                
         except Exception as e:
-            return f"Error in species identification: {str(e)}", []
+            logger.error(f"Species identification error: {e}")
+            return f"Species identification failed: {str(e)}", []
     
-    def get_location_info(self, latitude, longitude):
-        """Get location information from coordinates"""
-        if latitude is None or longitude is None:
-            return "Location not provided", None
-            
+    def get_wikipedia_info(self, species_name):
+        """Get Wikipedia information with better error handling"""
         try:
-            location = self.geolocator.reverse(f"{latitude}, {longitude}")
+            # Clean species name
+            clean_name = species_name.split(',')[0].split('(')[0].strip()
             
-            # Create a map
-            m = folium.Map(location=[latitude, longitude], zoom_start=15)
-            folium.Marker(
-                [latitude, longitude],
-                popup=f"Tree Location<br>Lat: {latitude:.6f}<br>Lon: {longitude:.6f}",
-                tooltip="Tree Location"
-            ).add_to(m)
+            search_queries = [
+                clean_name,
+                f"{clean_name} tree",
+                f"{clean_name} plant",
+                f"{clean_name} species"
+            ]
             
-            # Save map to temporary file
-            map_file = tempfile.NamedTemporaryFile(delete=False, suffix='.html')
-            m.save(map_file.name)
-            
-            address = location.address if location else "Address not found"
+            for query in search_queries:
+                try:
+                    results = wikipedia.search(query, results=2)
+                    if results:
+                        for result in results:
+                            try:
+                                page = wikipedia.page(result, auto_suggest=False)
+                                summary = wikipedia.summary(result, sentences=2, auto_suggest=False)
+                                return {
+                                    'title': page.title,
+                                    'summary': summary,
+                                    'url': page.url
+                                }
+                            except:
+                                continue
+                except:
+                    continue
             
-            return f"Location: {address}", map_file.name
+            return {
+                'title': 'No information found',
+                'summary': f'Wikipedia information not available for {species_name}',
+                'url': None
+            }
             
         except Exception as e:
-            return f"Error getting location info: {str(e)}", None
+            return {
+                'title': 'Error',
+                'summary': f'Could not retrieve information: {str(e)}',
+                'url': None
+            }
 
 def analyze_tree(image, latitude, longitude):
-    """Main function to analyze tree from image and location"""
+    """Main analysis function"""
     if image is None:
-        return "Please upload an image", "", "", "", None
-    
-    analyzer = TreeAnalyzer()
-    
-    # Analyze height
-    height_result, depth_map = analyzer.estimate_tree_height(image)
-    
-    # Identify species
-    species_result, species_info = analyzer.identify_tree_species(image)
+        return "Please upload an image", "", "", "", ""
     
-    # Get location info
-    location_result, map_file = analyzer.get_location_info(latitude, longitude)
-    
-    # Format species information
-    species_text = ""
-    if species_info:
-        for info in species_info:
-            species_text += f"**{info['species']}** (Confidence: {info['confidence']:.2f})\n"
-            species_text += f"*{info['summary']}*\n"
-            if info['url']:
-                species_text += f"[Wikipedia Link]({info['url']})\n"
-            species_text += "\n"
-    
-    # Return results
-    return (
-        height_result,
-        species_result,
-        species_text,
-        location_result,
-        map_file
-    )
-
-# Create Gradio interface
-def create_interface():
-    with gr.Blocks(title="Tree Analysis App", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        # 🌳 Tree Analysis App
+    try:
+        analyzer = TreeAnalyzer()
         
-        Upload an image of a tree and optionally provide GPS coordinates to get:
-        - **Tree height estimation** using MiDaS depth estimation
-        - **Species identification** with Wikipedia information
-        - **Location mapping** of where the tree was captured
-        """)
+        # Height estimation
+        height_result = analyzer.estimate_tree_height(image)
         
-        with gr.Row():
-            with gr.Column(scale=1):
-                image_input = gr.Image(type="pil", label="Upload Tree Image")
-                
-                with gr.Row():
-                    latitude_input = gr.Number(
-                        label="Latitude (optional)", 
-                        placeholder="e.g., 40.7128",
-                        info="GPS latitude coordinate"
-                    )
-                    longitude_input = gr.Number(
-                        label="Longitude (optional)", 
-                        placeholder="e.g., -74.0060",
-                        info="GPS longitude coordinate"
-                    )
-                
-                analyze_btn = gr.Button("🔍 Analyze Tree", variant="primary")
-                
-            with gr.Column(scale=2):
-                with gr.Tab("Results"):
-                    height_output = gr.Textbox(label="Height Estimation", lines=2)
-                    species_status = gr.Textbox(label="Species Identification Status", lines=1)
-                    species_output = gr.Markdown(label="Species Information")
-                    location_output = gr.Textbox(label="Location Information", lines=2)
-                    
-                with gr.Tab("Location Map"):
-                    map_output = gr.HTML(label="Location Map")
+        # Species identification
+        species_status, species_info = analyzer.identify_tree_species(image)
         
-        # Connect the analyze button
-        analyze_btn.click(
-            fn=analyze_tree,
-            inputs=[image_input, latitude_input, longitude_input],
-            outputs=[height_output, species_status, species_output, location_output, map_output]
-        )
+        # Format species results
+        species_text = ""
+        if species_info:
+            for i, info in enumerate(species_info, 1):
+                species_text += f"## {i}. {info['species']}\n"
+                species_text += f"**Confidence:** {info['confidence']:.3f}\n"
+                species_text += f"**Plant-related:** {'Yes' if info['is_plant_related'] else 'Uncertain'}\n"
+                
+                wiki = info['wiki_info']
+                species_text += f"**Wikipedia:** {wiki['title']}\n"
+                species_text += f"{wiki['summary']}\n"
+                if wiki['url']:
+                    species_text += f"🔗 [Read more]({wiki['url']})\n"
+                species_text += "\n---\n"
+        else:
+            species_text = "No species information could be determined from this image."
         
-        # Example section
-        gr.Markdown("""
-        ## 📱 Usage Tips:
-        1. **Take a clear photo** of the tree with good lighting
-        2. **Include reference objects** (like people) for better height estimation
-        3. **Enable GPS** on your phone and note the coordinates
-        4. **Upload the image** and enter GPS coordinates if available
-        5. **Click Analyze** to get comprehensive tree information
+        # Location info
+        location_result = ""
+        map_html = ""
         
-        ## 🔗 Sharing:
-        This app generates a shareable link that you can send to others!
-        """)
-    
-    return demo
-
-# Main execution
-if __name__ == "__main__":
-    # Create and launch the interface
-    demo = create_interface()
-    
-    # Launch with sharing enabled to generate a public link
-    demo.launch(
-        share=True,  # This creates a shareable public link
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+        if latitude is not None and longitude is not None:
+            try:
+                location_result = f"Coordinates: {latitude:.6f}, {longitude:.6f}"
+                
+                # Create map
+                m = folium.Map(location=[latitude, longitude], zoom_start=15)
+                folium.Marker(
+                    [latitude, longitude],
+                    popup=f"Tree Location<br>{latitude:.6f}, {longitude:.6f}",
+                    tooltip="Tree Location"
+                ).add_to(m)
+                
+                # Save map
+                map_file = tempfile.NamedTemporaryFile(delete=False, suffix='.html', mode='w')
+                m.save(map_file.name)
+                map_file.close()
+                
+                with open(map_file.name, 'r', encoding='utf-8') as f:
+                    map_html = f.read()
+                os.unlink(map_file.name)
+                
+            except Exception as e:
+                location_result = f"Error processing location: {str(e)}"
+                map_html = "<p>Could not generate map</p>"
+        else:
+            location_result = "No GPS coordinates provided"
+            map_html = "<p>