Upload analyze_spatial_relationships.py

tools/analyze_spatial_relationships.py

@@ -0,0 +1,119 @@
+from smolagents import tool
+import os
+import json
+import math
+import cv2
+import numpy as np
+from typing import Optional
+
+@tool
+def analyze_spatial_relationships(image_path: Optional[str] = None, object_type1: Optional[str] = None, object_type2: Optional[str] = None) -> str:
+    """Analyze spatial relationships between objects in an image.
+    Args:
+        image_path: Path to the image file to analyze.
+        object_type1: First object type to analyze.
+        object_type2: Second object type to analyze.
+    Returns:
+        A description of spatial relationships between the objects.
+    """
+    try:
+        # Check if image path exists
+        if not image_path or not os.path.exists(image_path):
+            return f"Error: Image path '{image_path}' does not exist."
+
+        # Check if detection data exists, if not run detection
+        detection_data_path = f"{os.path.splitext(image_path)[0]}_detection_data.json"
+
+        if not os.path.exists(detection_data_path):
+            # Run detection first
+            from tools.detect_objects import detect_objects_with_roboflow
+            detect_objects_with_roboflow(image_path=image_path)
+
+            if not os.path.exists(detection_data_path):
+                return f"Error: Detection failed. Could not create detection data."
+
+        # Load existing detection data
+        with open(detection_data_path, "r") as f:
+            result = json.load(f)
+
+        # Load the original image to get dimensions
+        img = cv2.imread(image_path)
+        if img is None:
+            return f"Error: Could not read image at '{image_path}'."
+
+        img_height, img_width = img.shape[:2]
+
+        # Define "close" threshold as 30% of image width
+        close_threshold = img_width * 0.3
+
+        # Extract objects of the specified types
+        objects1 = []
+        objects2 = []
+
+        if "predictions" in result:
+            for prediction in result["predictions"]:
+                class_name = prediction["class"]
+
+                # Check if this prediction matches either object type
+                if object_type1 and object_type1.lower() in class_name.lower():
+                    if "x" in prediction and "y" in prediction:
+                        objects1.append({
+                            "class": class_name,
+                            "x": prediction["x"],
+                            "y": prediction["y"],
+                            "width": prediction["width"],
+                            "height": prediction["height"],
+                            "confidence": prediction["confidence"]
+                        })
+
+                if object_type2 and object_type2.lower() in class_name.lower():
+                    if "x" in prediction and "y" in prediction:
+                        objects2.append({
+                            "class": class_name,
+                            "x": prediction["x"],
+                            "y": prediction["y"],
+                            "width": prediction["width"],
+                            "height": prediction["height"],
+                            "confidence": prediction["confidence"]
+                        })
+
+        # If no objects found, return appropriate message
+        if not objects1:
+            return f"No objects matching '{object_type1}' were found in the image."
+
+        if not objects2:
+            return f"No objects matching '{object_type2}' were found in the image."
+
+        # Analyze spatial relationships
+        close_pairs = []
+
+        for obj1 in objects1:
+            for obj2 in objects2:
+                # Calculate Euclidean distance between centers
+                distance = math.sqrt(
+                    (obj1["x"] - obj2["x"])**2 +
+                    (obj1["y"] - obj2["y"])**2
+                )
+
+                # Check if they are close
+                if distance <= close_threshold:
+                    close_pairs.append({
+                        "object1": obj1["class"],
+                        "object2": obj2["class"],
+                        "distance": distance,
+                        "distance_percent": (distance / img_width) * 100
+                    })
+
+        # Format the results
+        if close_pairs:
+            output = f"Found {len(close_pairs)} instances where {object_type1} is close to {object_type2}:\n"
+            for i, pair in enumerate(close_pairs, 1):
+                output += f"{i}. {pair['object1']} is close to {pair['object2']} (distance: {pair['distance']:.1f} pixels, {pair['distance_percent']:.1f}% of image width)\n"
+
+            output += f"\nClose is defined as objects with centers within {close_threshold:.1f} pixels (30% of image width) from each other."
+            return output
+        else:
+            return f"No instances found where {object_type1} is close to {object_type2}. Close is defined as objects with centers within {close_threshold:.1f} pixels (30% of image width) from each other."
+
+    except Exception as e:
+        return f"Error analyzing spatial relationships: {str(e)}"