add utils for gemini pointing

Author: lf-zhao

predicators/spot_utils/perception/utils/gemini_parsing.py

@@ -0,0 +1,145 @@
+"""Utility functions for parsing Gemini model responses."""
+
+import json
+import logging
+from typing import Dict, List, Optional, Tuple, Union
+
+from PIL import Image
+
+
+def parse_gemini_point_response(
+    response: str,
+    image_size: Optional[Tuple[int, int]] = None,
+    normalize: bool = True,
+) -> Dict[str, List[Dict[str, Union[List[float], str]]]]:
+    """Parse Gemini point response to extract point coordinates and labels.
+    
+    Args:
+        response: Raw response from Gemini
+        image_size: Optional size of the input image (width, height) for normalization
+        normalize: Whether to normalize points to 0-1000 range
+        
+    Returns:
+        Dictionary with points and labels
+    """
+    try:
+        # Extract JSON from response
+        start = response.find("[")
+        end = response.rfind("]") + 1
+        json_str = response[start:end]
+        points_data = json.loads(json_str)
+        
+        # Validate and normalize points
+        normalized_points = []
+        for item in points_data:
+            point = item["point"]
+            if not isinstance(point, list) or len(point) != 2:
+                raise ValueError(f"Invalid point format: {point}")
+                
+            if normalize and image_size:
+                # Convert to normalized coordinates [0-1000]
+                y, x = point
+                y = y * 1000 / image_size[1]
+                x = x * 1000 / image_size[0]
+                point = [y, x]
+            
+            normalized_points.append({
+                "point": point,
+                "label": str(item["label"])
+            })
+        
+        return {"points": normalized_points}
+    except (json.JSONDecodeError, KeyError, ValueError) as e:
+        logging.error(f"Failed to parse Gemini point response: {e}")
+        logging.error(f"Response: {response}")
+        raise ValueError("Failed to parse Gemini point response") from e
+
+
+def parse_gemini_detection_response(
+    response: str,
+    image_size: Optional[Tuple[int, int]] = None,
+    normalize: bool = True,
+) -> Dict[str, List[Dict[str, Union[List[float], str]]]]:
+    """Parse Gemini detection response to extract bounding boxes and labels.
+    
+    Args:
+        response: Raw response from Gemini
+        image_size: Optional size of the input image (width, height) for normalization
+        normalize: Whether to normalize boxes to 0-1000 range
+        
+    Returns:
+        Dictionary with bounding boxes and labels
+    """
+    try:
+        # Extract JSON from response
+        start = response.find("[")
+        end = response.rfind("]") + 1
+        json_str = response[start:end]
+        detection_data = json.loads(json_str)
+        
+        # Validate and normalize boxes
+        normalized_boxes = []
+        for item in detection_data:
+            box = item["box_2d"]
+            if not isinstance(box, list) or len(box) != 4:
+                raise ValueError(f"Invalid box format: {box}")
+                
+            if normalize and image_size:
+                # Convert to normalized coordinates [0-1000]
+                x1, y1, x2, y2 = box
+                x1 = x1 * 1000 / image_size[0]
+                y1 = y1 * 1000 / image_size[1]
+                x2 = x2 * 1000 / image_size[0]
+                y2 = y2 * 1000 / image_size[1]
+                box = [x1, y1, x2, y2]
+            
+            normalized_boxes.append({
+                "box_2d": box,
+                "label": str(item["label"])
+            })
+        
+        return {"detections": normalized_boxes}
+    except (json.JSONDecodeError, KeyError, ValueError) as e:
+        logging.error(f"Failed to parse Gemini detection response: {e}")
+        logging.error(f"Response: {response}")
+        raise ValueError("Failed to parse Gemini detection response") from e
+
+
+def denormalize_point(
+    point: List[float],
+    image_size: Tuple[int, int],
+) -> List[float]:
+    """Convert normalized point [0-1000] back to image coordinates.
+    
+    Args:
+        point: Normalized point coordinates [y, x]
+        image_size: Size of the image (width, height)
+        
+    Returns:
+        Point coordinates in image space
+    """
+    y, x = point
+    y = y * image_size[1] / 1000
+    x = x * image_size[0] / 1000
+    return [y, x]
+
+
+def denormalize_box(
+    box: List[float],
+    image_size: Tuple[int, int],
+) -> List[float]:
+    """Convert normalized box [0-1000] back to image coordinates.
+    
+    Args:
+        box: Normalized box coordinates [x1, y1, x2, y2]
+        image_size: Size of the image (width, height)
+        
+    Returns:
+        Box coordinates in image space
+    """
+    x1, y1, x2, y2 = box
+    x1 = x1 * image_size[0] / 1000
+    y1 = y1 * image_size[1] / 1000
+    x2 = x2 * image_size[0] / 1000
+    y2 = y2 * image_size[1] / 1000
+    return [x1, y1, x2, y2] 

predicators/spot_utils/perception/utils/google_utils.py

@@ -0,0 +1,150 @@
+import os
+from typing import Any, List, Optional, Union
+
+import google.generativeai as genai
+from PIL import Image
+from tenacity import retry, stop_after_attempt, wait_random_exponential
+from google.generativeai.generative_models import GenerativeModel
+from google.generativeai.types import GenerationConfig
+
+
+# GEMINI_MODEL_NAME = "gemini-2.5.pro-exp-03-25"
+# GEMINI_MODEL_NAME = "gemini-2.5.pro"
+GEMINI_MODEL_NAME = "gemini-2.0-flash"
+
+
+class GeminiClient:
+    """Client for interacting with Gemini models."""
+
+    def __init__(self, api_key: Optional[str] = None):
+        """Initialize the Gemini client.
+
+        Args:
+            api_key: Optional API key. If not provided, will use GEMINI_API_KEY env var.
+        """
+        api_key = api_key or os.getenv("GEMINI_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "GEMINI_API_KEY environment variable must be set or api_key must be provided"
+            )
+        genai.configure(api_key=api_key)
+
+    def generate(
+        self,
+        model: str = GEMINI_MODEL_NAME,
+        prompt: str = "",
+        image: Optional[Union[str, Image.Image]] = None,
+        temperature: float = 0.0,
+        max_tokens: int = 4096,
+        **kwargs: Any
+    ) -> str:
+        """Generate a response from Gemini.
+
+        Args:
+            model: Name of the Gemini model to use
+            prompt: Text prompt
+            image: Optional image (path or PIL Image)
+            temperature: Sampling temperature
+            max_tokens: Maximum number of tokens to generate
+            **kwargs: Additional arguments for the model
+
+        Returns:
+            Generated text response
+        """
+        gen_model = genai.GenerativeModel(model)
+        
+        if image is None:
+            response = gen_model.generate_content(
+                prompt,
+                generation_config=GenerationConfig(
+                    temperature=temperature,
+                    max_output_tokens=max_tokens,
+                    **kwargs
+                ),
+            )
+        else:
+            if isinstance(image, str):
+                image = Image.open(image)
+            response = gen_model.generate_content(
+                [prompt, image],
+                generation_config=GenerationConfig(
+                    temperature=temperature,
+                    max_output_tokens=max_tokens,
+                    **kwargs
+                ),
+            )
+        
+        return response.text
+
+    def generate_stream(
+        self,
+        model: str = GEMINI_MODEL_NAME,
+        prompt: str = "",
+        image: Optional[Union[str, Image.Image]] = None,
+        temperature: float = 0.0,
+        max_tokens: int = 4096,
+        **kwargs: Any
+    ) -> List[str]:
+        """Generate a streaming response from Gemini.
+
+        Args:
+            model: Name of the Gemini model to use
+            prompt: Text prompt
+            image: Optional image (path or PIL Image)
+            temperature: Sampling temperature
+            max_tokens: Maximum number of tokens to generate
+            **kwargs: Additional arguments for the model
+
+        Returns:
+            List of text chunks
+        """
+        gen_model = genai.GenerativeModel(model)
+        
+        if image is None:
+            response = gen_model.generate_content(
+                prompt,
+                stream=True,
+                generation_config=GenerationConfig(
+                    temperature=temperature,
+                    max_output_tokens=max_tokens,
+                    **kwargs
+                ),
+            )
+        else:
+            if isinstance(image, str):
+                image = Image.open(image)
+            response = gen_model.generate_content(
+                [prompt, image],
+                stream=True,
+                generation_config=GenerationConfig(
+                    temperature=temperature,
+                    max_output_tokens=max_tokens,
+                    **kwargs
+                ),
+            )
+        
+        return [chunk.text for chunk in response]
+
+
+@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+def call_google_api(
+    message: Union[str, List[Union[Any, Image.Image]]],
+    model: str = "gemini-pro",  # gemini-pro, gemini-pro-vision
+) -> str:
+    try:
+        gen_model = genai.GenerativeModel(model)
+        response = gen_model.generate_content(message)
+        response.resolve()
+        return response.text
+    except Exception as e:
+        print(f"{type(e).__name__}: {e}")
+        raise e
+
+
+if __name__ == "__main__":
+    # Example usage
+    client = GeminiClient()
+    input = "What color are apples?"
+    print("input: {}".format(input))
+    output = call_google_api(input)
+    print("output: {}".format(output))