ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/interfaces/images/plugins/abstract.py

@@ -0,0 +1,58 @@
+from typing import Any
+from abc import ABC, abstractmethod
+from PIL import Image
+from navconfig.logging import logging
+
+
+class ImagePlugin(ABC):
+    """
+    ImagePlugin is a base class for image processing plugins.
+    It provides a common interface for image processing tasks.
+    Subclasses should implement the `analyze` method to define
+    the specific image processing logic.
+    """
+    column_name: str = "image_info"
+
+    def __init__(self, *args, **kwargs):
+        """
+        Initialize the ImagePlugin with an optional image path.
+
+        :param image: Path to the image file.
+        """
+        self.column_name = kwargs.get("column_name", self.column_name)
+        self.logger = logging.getLogger(__name__)
+
+    @abstractmethod
+    async def analyze(self, image: Image.Image, **kwargs) -> Any:
+        """
+        Analyze the image and perform the desired processing.
+
+        :param image: Image Bytes opened with PIL Image.open
+        """
+        raise NotImplementedError(
+            "Image Plugin must implement analyze() method."
+        )
+
+    async def __aenter__(self):
+        if hasattr(self, "open"):
+            await self.open()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if hasattr(self, "close"):
+            await self.close()
+        return True
+
+    async def start(self):
+        """
+        Start the plugin. This method can be overridden by subclasses
+        to perform any initialization or setup tasks.
+        """
+        pass
+        return self
+
+    async def dispose(self):
+        """
+        Dispose of the plugin resources.
+        """
+        return self

parrot/interfaces/images/plugins/analisys.py

@@ -0,0 +1,148 @@
+from pathlib import Path
+from typing import Union, Dict, List
+import yaml
+from pydantic import BaseModel
+from PIL import Image
+from .abstract import ImagePlugin
+from ....clients.google import GoogleModel, GoogleGenAIClient
+
+
+class AnalysisPlugin(ImagePlugin):
+    """Plugin for analyzing images."""
+    column_name: str = "image_features"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._model_name: str = kwargs.get(
+            "model_name", GoogleModel.GEMINI_2_5_FLASH.value
+        )
+        self.prompt_path = kwargs.get("prompt_path", Path.cwd() / "prompts")
+        self.prompt_file: Union[str, Path] = kwargs.get(
+            "prompt_file", "default_analysis_prompt.json"
+        )
+        self._structured_outputs: Dict[str, BaseModel] = kwargs.get(
+            "structured_outputs", {}
+        )
+        self.filter_by: List[str] = kwargs.get(
+            "filter_by", ["Boxes on Floor"]
+        )
+        self.filter_column: str = kwargs.get(
+            "filter_column", "category"
+        )
+
+    def _load_model(self, model_name: str) -> BaseModel:
+        """ Load the classification or categorization model based on the provided model name.
+        This method uses importlib to dynamically import the model class.
+        """
+        try:
+            module_path, class_name = model_name.rsplit('.', 1)
+            module = __import__(module_path, fromlist=[class_name])
+            return getattr(module, class_name)
+        except (ImportError, AttributeError) as e:
+            raise ValueError(
+                f"Failed to load categorization model: {model_name}. Error: {e}"
+            )
+
+    async def start(self, **kwargs):
+        """Initialize the plugin and load the prompt."""
+        if isinstance(self.prompt_file, str):
+            self.prompt_file = Path(self.prompt_file)
+        if not self.prompt_file.is_absolute():
+            self.prompt_file = self.prompt_path.joinpath(self.prompt_file)
+        if not self.prompt_file.exists():
+            raise FileNotFoundError(
+                f"Prompt file {self.prompt_file} does not exist."
+            )
+        # From the prompt File, load the prompt content:
+        content = self.prompt_file.read_text()
+        # open from YAML file:
+        try:
+            self.prompt = yaml.safe_load(content).get('analysis_configs', {})
+        except Exception as e:
+            raise ValueError(
+                f"Failed to decode YAML from prompt file {self.prompt_file}. Error: {e}"
+            )
+        # Iterate over all prompts and load the models:
+        for key, value in self.prompt.items():
+            structured_output = value.get("structured_output", None)
+            if structured_output is None:
+                self.logger.warning(
+                    f"No structured output defined for {key} in prompt file."
+                )
+                continue
+            try:
+                self._structured_outputs[key] = self._load_model(structured_output)
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to load structured output model for {key}. Error: {e}"
+                )
+
+    def _extract_analysis_results(self, result) -> dict:
+        """
+        Extract analysis results from AIMessage object and return as dictionary.
+
+        Args:
+            result: AIMessage object with structured_output containing InkWallAnalysis
+
+        Returns:
+            dict: Dictionary with analysis data
+        """
+        if hasattr(result, 'output') and result.output:
+            analysis_result = result.output
+            # Convert Pydantic model to dictionary
+            return analysis_result.dict()
+        else:
+            raise ValueError("No output found in the result object.")
+
+    async def analyze(self, image: Union[Path, Image.Image], **kwargs) -> dict:
+        """
+        Analyze the ink wall image and perform structured analysis.
+
+        :param image: Image Bytes opened with PIL Image.open
+        :return: A dictionary containing the updated detections column.
+        """
+        row = kwargs.get("row", None)
+        # Check filter condition
+        if hasattr(self, 'filter_column') and hasattr(self, 'filter_by'):
+            if row[self.filter_column] not in self.filter_by:
+                self.logger.info(
+                    f"Skipping for analysis for row {row.name} with category {row[self.filter_column]}"
+                )
+                return None
+        image_classification = row[self.filter_column]
+        if not image_classification:
+            self.logger.warning(
+                f"Row {row.name} has no valid category for analysis."
+            )
+            return None
+        structured_output = self._structured_outputs.get(image_classification, None)
+        if structured_output is None:
+            self.logger.error(
+                f"No structured output defined for image classification: {image_classification}"
+            )
+            return None
+        # Perform analysis based on the image classification
+        try:
+            async with GoogleGenAIClient() as client:
+                _result = await client.ask_to_image(
+                    image=image,
+                    prompt=self.prompt[image_classification]['prompt'],
+                    structured_output=structured_output,
+                    model=self._model_name
+                )
+                if _result and isinstance(_result.output, structured_output):
+                    content = self._extract_analysis_results(_result)
+                    return {
+                        "analysis:": image_classification,
+                        **content
+                    }
+                else:
+                    self.logger.error(
+                        f"Unexpected output format for {image_classification} analysis."
+                    )
+                    return None
+        except Exception as e:
+            self.logger.error(
+                f"Error during analysis for {image_classification}: {e}"
+            )
+            return None

parrot/interfaces/images/plugins/classify.py

@@ -0,0 +1,150 @@
+from pathlib import Path
+from typing import List, Union, Optional
+from enum import Enum, EnumMeta
+from pydantic import BaseModel, Field
+from PIL import Image
+from .abstract import ImagePlugin
+from ....clients.google import GoogleModel, GoogleGenAIClient
+
+DEFAULT_PROMPT = """
+You are an expert in retail image analysis. Your task is to classify the provided image into one of the following categories.
+Please read the definitions carefully and choose the single best fit.
+"""
+
+def is_model_class(cls) -> bool:
+    return isinstance(cls, type) and issubclass(cls, BaseModel)
+
+
+def is_enum_class(cls) -> bool:
+    return isinstance(cls, type) and issubclass(cls, Enum)
+
+class ImageCategory(str, Enum):
+    """Enumeration for retail image categories."""
+    INK_WALL = "Ink Wall"
+    FRONT_OF_STORE = "Front of Store"
+    SHELVES_WITH_PRODUCTS = "Shelves with Products"
+    BOXES_ON_FLOOR = "Boxes on Floor"
+    MERCHANDISING_ENDCAP = "Merchandising Endcap"
+    OTHER = "Other"
+
+
+class ImageClassification(BaseModel):
+    """Schema for classifying a retail image."""
+    category: ImageCategory = Field(
+        ...,
+        description="The best-fitting category for the image based on the provided definitions."
+    )
+    confidence_score: float = Field(
+        ..., ge=0.0, le=1.0,
+        description="The model's confidence in its classification, from 0.0 to 1.0."
+    )
+    reasoning: str = Field(
+        ...,
+        description="A brief explanation for why the image was assigned to this category."
+    )
+
+
+class ClassificationPlugin(ImagePlugin):
+    """
+    ClassificationPlugin is a plugin for performing image classification.
+    Uses Gemini 2.5 multimodal model for image classification tasks.
+    """
+    column_name: str = "image_classifications"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._model_name: str = kwargs.get(
+            "model_name", GoogleModel.GEMINI_2_5_FLASH.value
+        )
+        self.prompt: List[str] = kwargs.get("prompt", DEFAULT_PROMPT)
+        self.confidence: float = kwargs.get("confidence", 0.5)
+        self._classification_model = kwargs.get(
+            "classification_model", self._load_classification_model(
+                ImageClassification
+            )
+        )
+        self._category_model = kwargs.get(
+            "category_model", self._load_category_model(
+                ImageCategory
+            )
+        )
+
+    def _load_model(self, model_name: str) -> BaseModel:
+        """ Load the classification or categorization model based on the provided model name.
+        This method uses importlib to dynamically import the model class.
+        """
+        try:
+            module_path, class_name = model_name.rsplit('.', 1)
+            module = __import__(module_path, fromlist=[class_name])
+            return getattr(module, class_name)
+        except (ImportError, AttributeError) as e:
+            raise ValueError(
+                f"Failed to load categorization model: {model_name}. Error: {e}"
+            )
+
+    def _load_category_model(self, model_name: Union[str, Enum]) -> Enum:
+        """
+        Load the categorization model based on the provided model name.
+
+        Category Model is a BaseModel that defines the structure of the categorization result.
+        model_name will be the python path where the model is stored.
+        for example, resources.models.categorization_models.ImageCategory
+        uses importlib to dynamically import the model class.
+        """
+        if is_enum_class(model_name):
+            # Already a Enum instance, return it directly
+            return model_name
+        elif isinstance(model_name, str):
+            # Attempt to import the model class dynamically
+            return self._load_model(model_name)
+        else:
+            raise ValueError(
+                "Category model_name must be a string or a Enum instance."
+            )
+
+    def _load_classification_model(self, model_name: Union[str, BaseModel]) -> BaseModel:
+        """
+        Load the classification model based on the provided model name.
+        """
+        if is_model_class(model_name):
+            # Already a BaseModel instance, return it directly
+            return model_name
+        elif isinstance(model_name, str):
+            # Attempt to import the model class dynamically
+            return self._load_model(model_name)
+        else:
+            raise ValueError(
+                "Classification model_name must be a string or a BaseModel instance."
+            )
+
+    async def analyze(self, image: Union[Path, Image.Image], **kwargs) -> dict:
+        """
+        Analyze the image and classify it into a retail category.
+
+        :param image: Image Bytes opened with PIL Image.open
+        :return: A dictionary containing the classification result.
+        """
+        async with GoogleGenAIClient() as client:
+            _result = await client.ask_to_image(
+                image=image,
+                prompt=self.prompt,
+                structured_output=self._classification_model,
+                model=self._model_name
+            )
+            if _result and isinstance(_result.output, self._classification_model):
+                result = _result.output
+                # evaluate if confidence_score is above the threshold
+                if result.confidence_score < self.confidence:
+                    self.logger.warning(
+                        f"Classification confidence {result.confidence_score} "
+                        f"is below the threshold {self.confidence}. "
+                        "Returning None."
+                    )
+                    return None
+                # If the model returns a valid classification result
+                return result.dict()
+            else:
+                self.logger.error(
+                    "The model did not return a valid classification result."
+                )
+                return None

parrot/interfaces/images/plugins/classifybase.py

@@ -0,0 +1,182 @@
+from pathlib import Path
+from typing import List, Optional, Union
+from pydantic import BaseModel
+import pandas as pd
+from .abstract import ImagePlugin
+from ....clients.google import GoogleModel
+from ....models import ObjectDetectionResult
+
+
+def is_model_class(cls) -> bool:
+    return isinstance(cls, type) and issubclass(cls, BaseModel)
+
+
+DEFAULT_PROMPT = ''
+
+
+class ClassifyBase(ImagePlugin):
+    """
+    ClassifyBase is an Abstract base class for performing image classification.
+    Uses Gemini 2.5 multimodal model for image classification tasks.
+    """
+    column_name: str = "detections"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._model_name: str = kwargs.get(
+            "model_name", GoogleModel.GEMINI_2_5_FLASH.value
+        )
+        model = kwargs.get(
+            "detection_model",
+            ObjectDetectionResult
+        )
+        self.reference_image: Optional[Path] = kwargs.get("reference_image", None)
+        self._detection_model: Optional[BaseModel] = self._load_model(model)
+        self.prompt: List[str] = kwargs.get("prompt", DEFAULT_PROMPT)
+        self.filter_by: List[str] = kwargs.get(
+            "filter_by", ["Boxes on Floor"]
+        )
+        # Modified: filter_column can be None to disable filtering
+        self.filter_column: Optional[str] = kwargs.get("filter_column", None)
+
+    async def start(self, **kwargs):
+        if isinstance(self.reference_image, str):
+            self.reference_image = Path(self.reference_image)
+        if self.reference_image and not self.reference_image.is_absolute():
+            self.reference_image = Path.cwd() / self.reference_image
+        if self.reference_image and not self.reference_image.exists():
+            self.logger.warning(
+                f"Reference image {self.reference_image} does not exist. "
+                "Classification may not work as expected."
+            )
+            self.reference_image = None
+
+    def _load_model(self, model_name: str) -> BaseModel:
+        """ Load the classification or categorization model based on the provided model name.
+        This method uses importlib to dynamically import the model class.
+        """
+        if is_model_class(model_name):
+            # Already a BaseModel instance, return it directly
+            return model_name
+        try:
+            module_path, class_name = model_name.rsplit('.', 1)
+            module = __import__(module_path, fromlist=[class_name])
+            return getattr(module, class_name)
+        except (ImportError, AttributeError) as e:
+            raise ValueError(
+                f"Failed to load categorization model: {model_name}. Error: {e}"
+            )
+
+    def _is_valid_filter_value(self, value):
+        """Check if a filter value is valid (not NA/NaN/None)."""
+        if pd.isna(value):
+            return False
+        if value is None:
+            return False
+        return True
+
+    def _should_apply_filter(self) -> bool:
+        """
+        Determine if filtering should be applied based on filter_column.
+        Returns False if filter_column is None, True otherwise.
+        """
+        return self.filter_column is not None
+
+    def _filter_dataset(self, dataset: pd.DataFrame) -> pd.DataFrame:
+        """
+        Apply filtering to the dataset if filter_column is specified.
+        If filter_column is None, return the entire dataset unchanged.
+
+        Args:
+            dataset: Input DataFrame to potentially filter
+
+        Returns:
+            Filtered DataFrame or original dataset if no filtering should be applied
+        """
+        if not self._should_apply_filter():
+            self.logger.debug(
+                "Filter column is None - processing entire dataset without filtering"
+            )
+            return dataset
+
+        if self.filter_column not in dataset.columns:
+            self.logger.warning(
+                f"Filter column '{self.filter_column}' not found in dataset. "
+                "Processing entire dataset."
+            )
+            return dataset
+
+        # Apply filtering logic
+        if not self.filter_by:
+            self.logger.warning("filter_by is empty - processing entire dataset")
+            return dataset
+
+        # Filter the dataset based on filter_by values in filter_column
+        mask = dataset[self.filter_column].apply(
+            lambda x: self._is_valid_filter_value(x) and x in self.filter_by
+        )
+
+        filtered_dataset = dataset[mask].copy()
+
+        self.logger.info(
+            f"Filtered dataset from {len(dataset)} to {len(filtered_dataset)} rows "
+            f"using filter_column='{self.filter_column}' with values {self.filter_by}"
+        )
+
+        return filtered_dataset
+
+    async def process_dataset(self, dataset: pd.DataFrame) -> pd.DataFrame:
+        """
+        Process the dataset with optional filtering.
+        This method should be implemented by subclasses to handle the actual classification.
+
+        Args:
+            dataset: Input DataFrame
+
+        Returns:
+            Processed DataFrame with classification results
+        """
+        # Apply filtering (or not, depending on filter_column)
+        filtered_dataset = self._filter_dataset(dataset)
+
+        # Placeholder for actual classification logic
+        # Subclasses should override this method
+        processed_dataset = await self._classify_images(filtered_dataset)
+
+        return processed_dataset
+
+    async def _classify_images(self, dataset: pd.DataFrame) -> pd.DataFrame:
+        """
+        Perform the actual image classification.
+        This method should be implemented by subclasses.
+
+        Args:
+            dataset: Filtered (or unfiltered) DataFrame to classify
+
+        Returns:
+            DataFrame with classification results
+        """
+        raise NotImplementedError("Subclasses must implement _classify_images method")
+
+    def configure_filtering(
+        self,
+        filter_column: Optional[str] = None,
+        filter_by: Optional[List[str]] = None
+    ) -> None:
+        """
+        Dynamically configure filtering parameters.
+
+        Args:
+            filter_column: Column to filter by. Set to None to disable filtering.
+            filter_by: Values to filter by. Only used if filter_column is not None.
+        """
+        if filter_column is not None:
+            self.filter_column = filter_column
+
+        if filter_by is not None:
+            self.filter_by = filter_by
+
+        self.logger.info(
+            f"Filtering configured: filter_column='{self.filter_column}', "
+            f"filter_by={self.filter_by if self.filter_column else 'N/A (filtering disabled)'}"
+        )

parrot/interfaces/images/plugins/detect.py

@@ -0,0 +1,150 @@
+from pathlib import Path
+from typing import Union
+from pydantic import BaseModel
+from PIL import Image
+from ....clients.google import GoogleGenAIClient
+from .classifybase import ClassifyBase
+
+
+def is_model_class(cls) -> bool:
+    return isinstance(cls, type) and issubclass(cls, BaseModel)
+
+
+DEFAULT_PROMPT = """
+Analyze this retail image to identify Epson and competitor products.
+
+TARGET BRANDS: Epson, HP, Canon, Brother, Lexmark, Xerox, Ricoh, Kyocera, Sharp
+
+TASK:
+1. Count ALL visible products (boxed and unboxed)
+2. Identify each product's brand, type, and model (if visible)
+3. Provide approximate location as normalized coordinates (0.0 to 1.0)
+
+IMPORTANT:
+- Focus on printers, scanners, ink cartridges, toner, and related products
+- Include products that are partially visible
+- Use confidence scores (0.0 to 1.0) for each detection
+- Provide bounding boxes as [x1, y1, x2, y2] where (0,0) is top-left, (1,1) is bottom-right
+
+Respond with a complete JSON object following the specified schema.
+"""  # noqa
+
+
+class DetectionPlugin(ClassifyBase):
+    """
+    DetectionPlugin is a plugin for performing image detection.
+    Uses Gemini 2.5 multimodal model for image detection tasks.
+    """
+    column_name: str = "detections"
+    section_name: str = "detections"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.section_name = kwargs.get("section_name", self.section_name)
+
+    def _extract_detection_results(self, result) -> dict:
+        """
+        Extract detection results from AIMessage object and return as dictionary.
+
+        Args:
+            result: AIMessage object with structured_output containing ObjectDetectionResult
+
+        Returns:
+            dict: Dictionary with 'analysis' and 'detections' keys
+        """
+
+        # Check if we have structured output
+        print(hasattr(result, 'structured_output'), hasattr(result, 'output'))
+        if hasattr(result, 'output') and result.output:
+            detection_result = result.output
+
+            # Convert BoundingBox objects to dictionaries
+            detections = []
+            for bbox in detection_result.detections:
+                detection_dict = {
+                    "object_id": bbox.object_id,
+                    "brand": bbox.brand,
+                    "model": bbox.model,
+                    "product_type": bbox.product_type,
+                    "description": bbox.description,
+                    "confidence": bbox.confidence,
+                    "bbox": bbox.bbox  # Already a list [x1, y1, x2, y2]
+                }
+                detections.append(detection_dict)
+
+            # Create the final result dictionary
+            detected = False
+            if detection_result.total_count > 0:
+                detected = True
+            r = {
+                "analysis": detection_result.analysis,
+                "total_count": detection_result.total_count,  # Include this for completeness
+                "detected": detected,
+                "detections": detections
+            }
+
+            return r
+
+        else:
+            # Fallback if no structured output
+            return {
+                "analysis": "No structured output available",
+                "total_count": 0,
+                "detections": []
+            }
+
+    async def analyze(self, image: Union[Path, Image.Image], **kwargs) -> dict:
+        """
+        Analyze the image and classify it into a retail category.
+
+        :param image: Image Bytes opened with PIL Image.open
+        :return: A dictionary containing the classification result.
+        """
+        row = kwargs.get("row", None)
+        detections_column = kwargs.get(self.column_name, None)
+        if detections_column is None:
+            detections_column = {}
+        else:
+            # Make a copy to avoid modifying original
+            detections_column = detections_column.copy()
+        if self.section_name in detections_column:
+            del detections_column[self.section_name]
+
+        if getattr(self, 'filter_column', None) and getattr(self, 'filter_by', None):
+            filter_value = row[self.filter_column] if self.filter_column in row else None
+            # Check if filter value is valid and not NA
+            if not self._is_valid_filter_value(filter_value):
+                self.logger.info(
+                    f"Skipping detection for row {row.name} - filter column '{self.filter_column}' has NA/invalid value"
+                )
+                return detections_column
+            # Now safe to do the comparison
+            if filter_value not in self.filter_by:
+                self.logger.info(
+                    f"Skipping detection for row {row.name} with category {filter_value}"
+                )
+                return detections_column
+        # Open the image
+        async with GoogleGenAIClient() as client:
+            _result = await client.ask_to_image(
+                image=image,
+                reference_images=[self.reference_image] if self.reference_image else None,
+                prompt=self.prompt,
+                structured_output=self._detection_model,
+                model=self._model_name
+            )
+            if _result:
+                detections = self._extract_detection_results(_result)
+                self.logger.info(f"Successfully detected {detections['total_count']} products")
+                self.logger.debug(
+                    f"Analysis: {detections['analysis'][:100]}..."
+                )
+                detections_column[self.section_name] = detections
+                row[self.column_name] = detections_column
+                # Return the updated detections column
+                return detections_column
+            else:
+                self.logger.error(
+                    "The model did not return a valid Detection result."
+                )
+                return detections_column