mb-rag 1.0.117__py3-none-any.whl → 1.0.124__py3-none-any.whl

mb_rag/utils/bounding_box.py

@@ -1,96 +1,223 @@
-import os
-import importlib.util
-from mb_rag.chatbot.basic import get_chatbot_google_generative_ai
+"""
+Bounding box utilities
+"""
 
-__all__ = [ "generate_bounding_box", "add_bounding_box"]
+import os
+from typing import List, Dict, Any, Optional, Tuple, Union
+from dataclasses import dataclass
+from mb_rag.utils.extra import check_package
 
+__all__ = ['BoundingBoxConfig', 'BoundingBoxProcessor']
 
-def check_package(package_name):
+def check_image_dependencies() -> None:
     """
-    Check if a package is installed
-    Args:
-        package_name (str): Name of the package
-    Returns:
-        bool: True if package is installed, False otherwise
+    Check if required image processing packages are installed
+    Raises:
+        ImportError: If any required package is missing
     """
-    return importlib.util.find_spec(package_name) is not None
+    if not check_package("PIL"):
+        raise ImportError("Pillow package not found. Please install it using: pip install Pillow")
+    if not check_package("cv2"):
+        raise ImportError("OpenCV package not found. Please install it using: pip install opencv-python")
+    if not check_package("google.generativeai"):
+        raise ImportError("Google Generative AI package not found. Please install it using: pip install google-generativeai")
 
+@dataclass
+class BoundingBoxConfig:
+    """Configuration for bounding box operations"""
+    model_name: str = "gemini-1.5-pro-latest"
+    api_key: Optional[str] = None
+    default_prompt: str = 'Return bounding boxes of container, for each only one return [ymin, xmin, ymax, xmax]'
 
-def generate_bounding_box(image_path: str,model_name: str = "gemini-1.5-flash", prompt: str = 'Return bounding boxes of container, for each only one return [ymin, xmin, ymax, xmax]'):
+class BoundingBoxProcessor:
     """
-    Function to generate bounding boxes
-    Args:
-        image_path (str): Image path
-        model_name (GenerativeModel): GenerativeModel object - google model (Default: )
-        prompt (str): Prompt
-    Returns:
-        res (str): Result
-    Raises:
-        FileNotFoundError: If image file doesn't exist
-        ValueError: If model is None or invalid
-    """
-    try:
-        model = get_chatbot_google_generative_ai(model_name)
-    except Exception as e:
-        raise ValueError(f"Error initializing model: {str(e)}")
-    
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image file not found: {image_path}")
+    Class for processing images and generating bounding boxes
     
-    try:
-        if not check_package('PIL'):
-            raise ImportError("PIL package not found. Please install it using: pip install pillow")    
-        from PIL import Image     
+    Attributes:
+        model: The Google Generative AI model instance
+        config: Configuration for bounding box operations
+    """
 
-        image = Image.open(image_path)
-        res = model.generate_content([image, prompt])
-        return res
-    except Exception as e:
-        raise ValueError(f"Error generating bounding boxes: {str(e)}")
+    def __init__(self, config: Optional[BoundingBoxConfig] = None, **kwargs):
+        """
+        Initialize bounding box processor
+        Args:
+            config: Configuration for the processor
+            **kwargs: Additional arguments
+        """
+        check_image_dependencies()
+        self.config = config or BoundingBoxConfig(**kwargs)
+        self._initialize_model()
+        self._initialize_image_libs()
 
-def add_bounding_box(image_path: str, bounding_box: list, show: bool = False):
-    """
-    Function to add bounding box to image
-    Args:
-        image_path (str): Image path
-        bounding_box (list): Bounding boxes
-        show (bool): Whether to display the image
-    Returns:
-        image (Image): Image with bounding box
-    Raises:
-        FileNotFoundError: If image file doesn't exist
-        ValueError: If bounding box format is invalid
-    """
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image file not found: {image_path}")
-    
-    if not isinstance(bounding_box, list):
-        raise ValueError("bounding_box must be a list of bounding box coordinates with keys as labels and values as [ymin, xmin, ymax, xmax]")
-    
-    try:
-        if not check_package('cv2'):
-            raise ImportError("cv2 package not found. Please install it using: pip install opencv-python")    
+    @classmethod
+    def from_model(cls, model_name: str, api_key: Optional[str] = None, **kwargs) -> 'BoundingBoxProcessor':
+        """
+        Create processor with specific model configuration
+        Args:
+            model_name: Name of the model
+            api_key: Optional API key
+            **kwargs: Additional configuration
+        Returns:
+            BoundingBoxProcessor: Configured processor
+        """
+        config = BoundingBoxConfig(
+            model_name=model_name,
+            api_key=api_key
+        )
+        return cls(config, **kwargs)
+
+    def _initialize_model(self) -> None:
+        """Initialize the AI model"""
+        import google.generativeai as genai
+        
+        api_key = self.config.api_key or os.environ.get("GOOGLE_API_KEY")
+        if not api_key:
+            raise ValueError("Google API key not found. Please provide api_key parameter or set GOOGLE_API_KEY environment variable.")
+        
+        try:
+            genai.configure(api_key=api_key)
+            self.model = genai.GenerativeModel(model_name=self.config.model_name)
+        except Exception as e:
+            raise ValueError(f"Error initializing Google Generative AI model: {str(e)}")
+
+    def _initialize_image_libs(self) -> None:
+        """Initialize image processing libraries"""
+        from PIL import Image
         import cv2
-        img = cv2.imread(image_path)
-        if img is None:
-            raise ValueError(f"Failed to load image: {image_path}")
+        self._Image = Image
+        self._cv2 = cv2
+
+    @staticmethod
+    def _validate_image_path(image_path: str) -> None:
+        """
+        Validate image path
+        Args:
+            image_path: Path to image
+        Raises:
+            FileNotFoundError: If image doesn't exist
+        """
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    def generate_bounding_boxes(self, 
+                              image_path: str, 
+                              prompt: Optional[str] = None) -> Any:
+        """
+        Generate bounding boxes for an image
+        Args:
+            image_path: Path to image
+            prompt: Custom prompt for the model
+        Returns:
+            Any: Model response with bounding boxes
+        """
+        self._validate_image_path(image_path)
         
+        try:
+            image = self._Image.open(image_path)
+            prompt = prompt or self.config.default_prompt
+            return self.model.generate_content([image, prompt])
+        except Exception as e:
+            raise ValueError(f"Error generating bounding boxes: {str(e)}")
 
-        for box in bounding_box:
-            if not isinstance(box, dict):
-                raise ValueError("bounding_box must be a list of bounding box coordinates with keys as labels and values as [ymin, xmin, ymax, xmax]")
-            for key, value in box.items():
+    def add_bounding_boxes(self, 
+                          image_path: str,
+                          bounding_boxes: Dict[str, List[int]],
+                          color: Tuple[int, int, int] = (0, 0, 255),
+                          thickness: int = 4,
+                          font_scale: float = 1.0,
+                          show: bool = False) -> Any:
+        """
+        Add bounding boxes to an image
+        Args:
+            image_path: Path to image
+            bounding_boxes: Dictionary of bounding boxes
+            color: BGR color tuple
+            thickness: Line thickness
+            font_scale: Font scale for labels
+            show: Whether to display the image
+        Returns:
+            Any: Image with bounding boxes
+        """
+        self._validate_image_path(image_path)
+        
+        if not isinstance(bounding_boxes, dict):
+            raise ValueError("bounding_boxes must be a dictionary")
+        
+        try:
+            img = self._cv2.imread(image_path)
+            if img is None:
+                raise ValueError(f"Failed to load image: {image_path}")
+            
+            for key, value in bounding_boxes.items():
                 if not isinstance(value, list) or len(value) != 4:
                     raise ValueError(f"Invalid bounding box format for key {key}. Expected [ymin, xmin, ymax, xmax]")
+                
+                self._cv2.rectangle(
+                    img=img,
+                    pt1=(value[1], value[0]),  # xmin, ymin
+                    pt2=(value[3], value[2]),  # xmax, ymax
+                    color=color,
+                    thickness=thickness
+                )
+                self._cv2.putText(
+                    img=img,
+                    text=key,
+                    org=(value[1], value[0]),
+                    fontFace=self._cv2.FONT_HERSHEY_SIMPLEX,
+                    fontScale=font_scale,
+                    color=color,
+                    thickness=thickness//2
+                )
+            
+            if show:
+                self._display_image(img)
             
-                cv2.rectangle(img, (value[1], value[0]), (value[3], value[2]), (0, 0, 255), 4)
-                cv2.putText(img, key, (value[1], value[0]), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+            return img
+        except Exception as e:
+            raise ValueError(f"Error adding bounding boxes to image: {str(e)}")
+
+    def _display_image(self, img: Any) -> None:
+        """
+        Display an image
+        Args:
+            img: Image to display
+        """
+        self._cv2.imshow("Image", img)
+        self._cv2.waitKey(0)
+        self._cv2.destroyAllWindows()
+
+    def save_image(self, img: Any, output_path: str) -> None:
+        """
+        Save an image
+        Args:
+            img: Image to save
+            output_path: Path to save the image
+        """
+        try:
+            self._cv2.imwrite(output_path, img)
+        except Exception as e:
+            raise ValueError(f"Error saving image: {str(e)}")
+
+    def process_image(self, 
+                     image_path: str,
+                     output_path: Optional[str] = None,
+                     show: bool = False,
+                     **kwargs) -> Any:
+        """
+        Complete image processing pipeline
+        Args:
+            image_path: Path to input image
+            output_path: Optional path to save output
+            show: Whether to display the result
+            **kwargs: Additional arguments for bounding box generation
+        Returns:
+            Any: Processed image
+        """
+        boxes = self.generate_bounding_boxes(image_path, **kwargs)
+        img = self.add_bounding_boxes(image_path, boxes, show=show)
         
-        if show:
-            cv2.imshow("Image", img)
-            cv2.waitKey(0)
-            cv2.destroyAllWindows()
+        if output_path:
+            self.save_image(img, output_path)
         
         return img
-    except Exception as e:
-        raise ValueError(f"Error adding bounding box to image: {str(e)}")

mb_rag/utils/extra.py

@@ -4,7 +4,7 @@ import os
 from dotenv import load_dotenv
 import importlib.util
 
-__all__ = ["load_env_file"]
+__all__ = ["load_env_file", "check_package", "pdf_to_text", "convert_pdfs_in_folder"]
 
 def load_env_file(file_path='.env'):
     """

mb_rag/version.py

@@ -1,5 +1,5 @@
 MAJOR_VERSION = 1
 MINOR_VERSION = 0
-PATCH_VERSION = 117
+PATCH_VERSION = 124
 version = '{}.{}.{}'.format(MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION)
 __all__  = ['MAJOR_VERSION', 'MINOR_VERSION', 'PATCH_VERSION', 'version']

{mb_rag-1.0.117.dist-info → mb_rag-1.0.124.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mb_rag
-Version: 1.0.117
+Version: 1.0.124
 Summary: RAG function file
 Author: ['Malav Bateriwala']
 Requires-Python: >=3.8

mb_rag-1.0.124.dist-info/RECORD

@@ -0,0 +1,15 @@
+mb_rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mb_rag/version.py,sha256=g3pkzzTRM6lWK4_vY_dwd2hYFDmXrp2VRW8Uj2krk4k,208
+mb_rag/chatbot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mb_rag/chatbot/basic.py,sha256=OR2IvDg-Sy968C2Mna6lxmFfh7Czj8yCEkCfyvtxBwI,14223
+mb_rag/chatbot/chains.py,sha256=vDbLX5R29sWN1pcFqJ5fyxJEgMCM81JAikunAEvMC9A,7223
+mb_rag/chatbot/prompts.py,sha256=n1PyiLbU-5fkslRv6aVOzt0dDlwya_cEdQ7kRnRhMuY,1749
+mb_rag/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mb_rag/rag/embeddings.py,sha256=kOnHjrbi0GRVErfcjML8fZz-KttipObUfa5fW9tGOoY,21196
+mb_rag/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mb_rag/utils/bounding_box.py,sha256=XnuhnLrsGvsI8P8VtOwlBrDlFE2It1HEZOcLlK6kusE,7931
+mb_rag/utils/extra.py,sha256=spbFrGgdruNyYQ5PzgvpSIa6Nm0rn9bb4qc8W9g582o,2492
+mb_rag-1.0.124.dist-info/METADATA,sha256=XyQ055JYBEiv5OU8m9FBhcOnHCoipo69LbnMgSa4bmM,154
+mb_rag-1.0.124.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+mb_rag-1.0.124.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
+mb_rag-1.0.124.dist-info/RECORD,,

mb_rag-1.0.117.dist-info/RECORD

@@ -1,15 +0,0 @@
-mb_rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/version.py,sha256=gwipOtP0TZlqdg5Rg1j43hyxg2GzlRZ0U9N3VPUDxNE,208
-mb_rag/chatbot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/chatbot/basic.py,sha256=AEcgMrXcH-QsCdoykqqgRJsS2-ZOEmxJuj2nd5KgcgI,13804
-mb_rag/chatbot/chains.py,sha256=KwdvAgI3DosSLgXGsm6fxeOKPYSbwy1mHH4ZIlOfTZY,8355
-mb_rag/chatbot/prompts.py,sha256=n1PyiLbU-5fkslRv6aVOzt0dDlwya_cEdQ7kRnRhMuY,1749
-mb_rag/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/rag/embeddings.py,sha256=kOnHjrbi0GRVErfcjML8fZz-KttipObUfa5fW9tGOoY,21196
-mb_rag/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/utils/bounding_box.py,sha256=PLfP_wSFkfNRu1uOu-1YIAQh93cuUHIkpFj77M4GaNk,3676
-mb_rag/utils/extra.py,sha256=qgaFsybeonT6XAFNfAU5h24FimDb-beGZ1qTztTFhSk,2434
-mb_rag-1.0.117.dist-info/METADATA,sha256=YH6LUu87Vn06x6NLD8FNS1YnaOuqx8Tq1vpTb8_81W0,154
-mb_rag-1.0.117.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-mb_rag-1.0.117.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
-mb_rag-1.0.117.dist-info/RECORD,,