speedy-utils 1.1.26__py3-none-any.whl → 1.1.28__py3-none-any.whl

vision_utils/plot.py

@@ -0,0 +1,345 @@
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
+
+from speedy_utils.__imports import np, plt
+
+
+if TYPE_CHECKING:
+
+    import lazy_loader as lazy
+    import matplotlib.pyplot as plt
+    import numpy as np
+    import torch
+
+    from .io_utils import read_images
+
+
+def _check_torch_available():
+    """Check if torch is available without importing at module level."""
+    try:
+        import torch
+
+        return True, torch
+    except ImportError:
+        return False, None
+
+
+def _check_matplotlib_available():
+    """Check if matplotlib is available without importing at module level."""
+    try:
+        import matplotlib.pyplot as plt
+
+        return True, plt
+    except ImportError:
+        return False, None
+
+
+def _to_numpy(img: Any) -> np.ndarray:
+    """Convert image to numpy array."""
+    torch_available, torch = _check_torch_available()
+    if torch_available and torch is not None and isinstance(img, torch.Tensor):
+        return img.detach().cpu().numpy()
+    if isinstance(img, np.ndarray):
+        return img
+    raise TypeError(
+        f'Unsupported image type: {type(img)}. Expected numpy.ndarray or torch.Tensor'
+    )
+
+
+def _normalize_image_format(img: np.ndarray) -> np.ndarray:
+    """
+    Normalize image to (H, W, C) format.
+
+    Detects and converts from:
+    - (C, H, W) where C is 1 or 3
+    - (H, W) grayscale
+    - (H, W, C) already correct
+    """
+    if img.ndim == 2:
+        # Grayscale (H, W) -> (H, W, 1)
+        return img[:, :, np.newaxis]
+    if img.ndim == 3:
+        # Check if it's (C, H, W) format
+        if img.shape[0] in [1, 3] and img.shape[0] < min(img.shape[1:]):
+            # Likely (C, H, W) -> transpose to (H, W, C)
+            return np.transpose(img, (1, 2, 0))
+        # Already (H, W, C)
+        return img
+    raise ValueError(f'Invalid image shape: {img.shape}. Expected 2D or 3D array')
+
+
+def _normalize_batch(
+    images: Union[np.ndarray, List[np.ndarray], List[Any], Any],
+) -> List[np.ndarray]:
+    """
+    Normalize batch of images to list of (H, W, C) numpy arrays.
+
+    Handles:
+    - List of numpy arrays or torch tensors
+    - List of file paths (strings or Path objects)
+    - Single numpy array of shape (B, H, W, C) or (B, C, H, W)
+    - Single torch tensor of shape (B, H, W, C) or (B, C, H, W)
+    """
+    # Convert to numpy if torch tensor
+    torch_available, torch = _check_torch_available()
+    if torch_available and torch is not None and isinstance(images, torch.Tensor):
+        images = images.detach().cpu().numpy()
+
+    # Handle single numpy array with batch dimension
+    if isinstance(images, np.ndarray):
+        if images.ndim == 4:
+            # (B, H, W, C) or (B, C, H, W)
+            # Check if it's (B, C, H, W) format
+            if images.shape[1] in [1, 3] and images.shape[1] < min(images.shape[2:]):
+                # (B, C, H, W) -> transpose to (B, H, W, C)
+                images = np.transpose(images, (0, 2, 3, 1))
+            # Convert to list of images
+            images = [images[i] for i in range(images.shape[0])]
+        elif images.ndim == 3:
+            # Single image (H, W, C) or (C, H, W)
+            images = [_normalize_image_format(images)]
+        elif images.ndim == 2:
+            # Single grayscale image (H, W)
+            images = [images[:, :, np.newaxis]]
+        else:
+            raise ValueError(
+                f'Invalid array shape: {images.shape}. Expected 2D, 3D, or 4D array'
+            )
+
+    # Handle list of images
+    if isinstance(images, list):
+        path_indices = [
+            idx for idx, img in enumerate(images) if isinstance(img, (str, Path))
+        ]
+
+        # Bulk load any file paths while preserving order
+        loaded_paths = {}
+        if path_indices:
+            loaded_arrays = read_images([str(images[idx]) for idx in path_indices])
+
+            if len(loaded_arrays) != len(path_indices):
+                raise ValueError(
+                    'Number of loaded images does not match number of paths provided.'
+                )
+
+            for idx, arr in zip(path_indices, loaded_arrays, strict=False):
+                loaded_paths[idx] = arr
+
+        normalized = []
+        for idx, img in enumerate(images):
+            if idx in loaded_paths:
+                img_np = loaded_paths[idx]
+            else:
+                img_np = _to_numpy(img)
+            img_normalized = _normalize_image_format(img_np)
+            normalized.append(img_normalized)
+        return normalized
+
+    raise TypeError(
+        f'Unsupported images type: {type(images)}. '
+        'Expected list, numpy.ndarray, or torch.Tensor'
+    )
+
+
+def plot_images_notebook(
+    images: Union[np.ndarray, List[np.ndarray], List[Any], Any],
+    nrows: Optional[int] = None,
+    ncols: Optional[int] = None,
+    figsize: Optional[Tuple[float, float]] = None,
+    titles: Optional[List[str]] = None,
+    cmap: Optional[str] = None,
+    dpi: int = 300,
+    max_figure_width: float = 15.0,
+    max_figure_height: float = 20.0,
+):
+    """
+    Plot a batch of images in a notebook with smart grid layout.
+    Handles images of different shapes gracefully.
+
+    Args:
+        images: Images to plot. Can be:
+            - List of numpy arrays or torch tensors (can have different shapes)
+            - List of image file paths (strings or Path objects)
+            - Numpy array of shape (B, H, W, C) or (B, C, H, W)
+            - Torch tensor of shape (B, H, W, C) or (B, C, H, W)
+            Each image can be (H, W), (H, W, C), (C, H, W) format
+        nrows: Number of rows in grid. If None, auto-calculated from sqrt
+        ncols: Number of columns in grid. If None, auto-calculated from sqrt
+        figsize: Figure size (width, height). If None, auto-calculated
+        titles: List of titles for each image
+        cmap: Colormap for grayscale images (default: 'gray')
+        dpi: Dots per inch for the figure (default: 72)
+        max_figure_width: Maximum figure width in inches (default: 15)
+        max_figure_height: Maximum figure height in inches (default: 20)
+
+    Example:
+        >>> import numpy as np
+        >>> # Auto grid layout with sqrt
+        >>> images = np.random.rand(9, 64, 64, 3)
+        >>> plot_images_notebook(images)  # 3x3 grid
+
+        >>> # Custom grid
+        >>> images = np.random.rand(8, 64, 64, 3)
+        >>> plot_images_notebook(images, nrows=2, ncols=4)
+
+        >>> # PyTorch tensor in (B, C, H, W) format
+        >>> import torch
+        >>> images = torch.rand(8, 3, 64, 64)
+        >>> plot_images_notebook(images)
+
+        >>> # List of images with different formats and shapes
+        >>> images = [
+        ...     np.random.rand(64, 64, 3),    # (H, W, C)
+        ...     np.random.rand(3, 128, 128),  # (C, H, W) - different size
+        ...     torch.rand(32, 48),           # Grayscale - different size
+        ...     np.random.rand(100, 200, 3),  # Different aspect ratio
+        ... ]
+        >>> plot_images_notebook(images, ncols=2)
+    """
+    # Check matplotlib availability
+    mpl_available, plt = _check_matplotlib_available()
+    if not mpl_available:
+        raise ImportError("matplotlib is required for plotting. Install it with: pip install matplotlib")
+
+    # Normalize all images to list of (H, W, C) numpy arrays
+    images_list = _normalize_batch(images)
+
+    n_images = len(images_list)
+
+    # Smart grid layout calculation
+    if nrows is None and ncols is None:
+        # Use sqrt to get roughly square grid
+        ncols_calc = int(np.ceil(np.sqrt(n_images)))
+        nrows_calc = int(np.ceil(n_images / ncols_calc))
+        nrows = nrows_calc
+        ncols = ncols_calc
+    elif nrows is None:
+        # Calculate rows from columns
+        assert ncols is not None
+        nrows = int(np.ceil(n_images / ncols))
+    elif ncols is None:
+        # Calculate columns from rows
+        assert nrows is not None
+        ncols = int(np.ceil(n_images / nrows))
+
+    # At this point, both nrows and ncols are guaranteed to be int
+    assert nrows is not None and ncols is not None
+
+    # Auto-calculate figure size if not provided
+    if figsize is None:
+        # Calculate based on average aspect ratio across all images
+        avg_aspect_ratio = 0.0
+        for img in images_list:
+            img_height, img_width = img.shape[:2]
+            avg_aspect_ratio += img_width / img_height
+        avg_aspect_ratio /= n_images
+
+        # Target cell size in inches (smaller for many images)
+        if n_images <= 4:
+            cell_width = 4.0
+        elif n_images <= 9:
+            cell_width = 3.0
+        elif n_images <= 16:
+            cell_width = 2.5
+        else:
+            cell_width = 2.0
+
+        cell_height = cell_width / avg_aspect_ratio
+
+        fig_width = ncols * cell_width
+        fig_height = nrows * cell_height
+
+        # Constrain to max sizes to prevent notebook breaking
+        if fig_width > max_figure_width:
+            scale = max_figure_width / fig_width
+            fig_width = max_figure_width
+            fig_height *= scale
+
+        if fig_height > max_figure_height:
+            scale = max_figure_height / fig_height
+            fig_height = max_figure_height
+            fig_width *= scale
+
+        figsize = (fig_width, fig_height)
+
+    fig, axes = plt.subplots(nrows, ncols, figsize=figsize, dpi=dpi, squeeze=False)
+
+    # Flatten axes for easier iteration
+    axes_flat = axes.flatten()
+
+    for idx, (ax, img) in enumerate(zip(axes_flat, images_list, strict=False)):
+        # Determine if grayscale
+        is_grayscale = img.shape[-1] == 1
+
+        if is_grayscale:
+            ax.imshow(img[:, :, 0], cmap=cmap or 'gray', aspect='auto')
+        else:
+            # Clip values to [0, 1] if they look like normalized images
+            if img.max() <= 1.0:
+                img_display = np.clip(img, 0, 1)
+            else:
+                # Assume [0, 255] range
+                img_display = np.clip(img / 255.0, 0, 1)
+            ax.imshow(img_display, aspect='auto')
+
+        ax.axis('off')
+
+        if titles and idx < len(titles):
+            ax.set_title(titles[idx], fontsize=8 if n_images > 9 else 10)
+
+    # Hide unused subplots
+    for idx in range(n_images, len(axes_flat)):
+        axes_flat[idx].axis('off')
+
+    plt.tight_layout()
+    plt.show()
+
+
+def visualize_tensor(img_tensor, mode='hwc', normalize=True, max_cols=8):
+    """
+    Visualize a tensor as an image or grid.
+
+    Args:
+        img_tensor: torch.Tensor, shape (C,H,W), (H,W,C), or (B,C,H,W)
+        mode: "hwc", "chw", or "bchw"
+        normalize: scale float tensor to 0–255 uint8 for display
+        max_cols: max columns when tiling a batch
+    """
+    # Check matplotlib availability
+    mpl_available, plt = _check_matplotlib_available()
+    if not mpl_available:
+        raise ImportError("matplotlib is required for plotting. Install it with: pip install matplotlib")
+    
+    if mode == 'chw':
+        img_tensor = img_tensor.permute(1, 2, 0)
+        imgs = [img_tensor]
+    elif mode == 'bchw':
+        b, c, h, w = img_tensor.shape
+        imgs = [img_tensor[i].permute(1, 2, 0) for i in range(b)]
+    elif mode == 'hwc':
+        imgs = [img_tensor]
+    else:
+        raise ValueError("mode must be 'hwc', 'chw', or 'bchw'")
+
+    # normalize each image
+    processed = []
+    for img in imgs:
+        img = img.detach().cpu().numpy()
+        if normalize:
+            img = (img - img.min()) / (img.max() - img.min() + 1e-8)
+            img = (img * 255).astype(np.uint8)
+        processed.append(img)
+
+    if len(processed) == 1:
+        plt.imshow(processed[0])
+    else:
+        cols = min(max_cols, len(processed))
+        rows = int(np.ceil(len(processed) / cols))
+        fig, axes = plt.subplots(rows, cols, figsize=(cols * 2, rows * 2))
+        axes = np.atleast_2d(axes)
+        for ax, img in zip(axes.flat, processed, strict=False):
+            ax.imshow(img)
+            ax.axis('off')
+        for ax in axes.flat[len(processed) :]:
+            ax.axis('off')
+    plt.show()

speedy_utils-1.1.26.dist-info/RECORD

@@ -1,52 +0,0 @@
-llm_utils/__init__.py,sha256=pbnOQddU5KnhP8uqMqN9E87BeDeCxFrgta2m2P89LmM,1591
-llm_utils/group_messages.py,sha256=Oe2tlhg-zRodG1-hodYebddrR77j9UdE05LzJw0EvYI,3622
-llm_utils/chat_format/__init__.py,sha256=MCNT8o-BZWmoOFE5VLyhJJOqHg8lJGqHXEKSXU08fK0,775
-llm_utils/chat_format/display.py,sha256=HiAOAC8FY7956gNuwE7rxii1MCCebn0avbXi1iIcDSc,17178
-llm_utils/chat_format/transform.py,sha256=eU0c3PdAHCNLuGP1UqPwln0B34Lv3bt_uV9v9BrlCN4,5402
-llm_utils/chat_format/utils.py,sha256=xTxN4HrLHcRO2PfCTR43nH1M5zCa7v0kTTdzAcGkZg0,1229
-llm_utils/lm/__init__.py,sha256=FBe8wVNWDMpvJ2kQYedJ3HH5L2BCAZBQVE0zEjND0Vo,729
-llm_utils/lm/base_prompt_builder.py,sha256=OLqyxbA8QeYIVFzB9EqxUiE_P2p4_MD_Lq4WSwxFtKU,12136
-llm_utils/lm/llm.py,sha256=uk45JhVcWDMaqezn9Yn_K5hehFSmQ4txU901fn_PcQg,16262
-llm_utils/lm/llm_signature.py,sha256=SP72cWXaVGcZs3m2V361DcLk_St7aYJamNapUiFBB6Q,1242
-llm_utils/lm/lm_base.py,sha256=pqbHZOdR7yUMpvwt8uBG1dZnt76SY_Wk8BkXQQ-mpWs,9557
-llm_utils/lm/mixins.py,sha256=Sn5KyPKGCT_HVJmmosmy3XSlZ0_k5Kds0VvSJqeUDpI,13695
-llm_utils/lm/openai_memoize.py,sha256=PDs3YCXKgHXaHlegkhouzPtf2Gom_o7pvzChCT-NQyQ,3870
-llm_utils/lm/signature.py,sha256=16QOHnGc-p7H8rR3j1dPg8AokdV_rEGUYCGGkIHIghE,10240
-llm_utils/lm/utils.py,sha256=oiJ50b8WV6oktnW4BByr1gRaGc55VJeF3IyhHqoofp4,12193
-llm_utils/lm/async_lm/__init__.py,sha256=PUBbCuf5u6-0GBUu-2PI6YAguzsyXj-LPkU6vccqT6E,121
-llm_utils/lm/async_lm/_utils.py,sha256=P1-pUDf_0pDmo8WTIi43t5ARlyGA1RIJfpAhz-gfA5g,6105
-llm_utils/lm/async_lm/async_llm_task.py,sha256=-BVOk18ZD8eC2obTLgiPq39f2PP3cji17Ku-Gb7c7Xo,18683
-llm_utils/lm/async_lm/async_lm.py,sha256=e3o9cyMbkVz_jQDTjJv2ybET_5mY012zdZGjNwi4Qk4,13719
-llm_utils/lm/async_lm/async_lm_base.py,sha256=iJgtzI6pVJzWtlXGqVLwgCIb-FzZAa3E5xW8yhyHUmM,8426
-llm_utils/lm/async_lm/lm_specific.py,sha256=KmqdCm3SJ5MqN-dRJd6S5tq5-ve1X2eNWf2CMFtc_3s,3926
-llm_utils/scripts/README.md,sha256=yuOLnLa2od2jp4wVy3rV0rESeiV3o8zol5MNMsZx0DY,999
-llm_utils/scripts/vllm_load_balancer.py,sha256=TT5Ypq7gUcl52gRFp--ORFFjzhfGlcaX2rkRv8NxlxU,37259
-llm_utils/scripts/vllm_serve.py,sha256=gJ0-y4kybMfSt8qzye1pJqGMY3x9JLRi6Tu7RjJMnss,14771
-llm_utils/vector_cache/__init__.py,sha256=i1KQuC4OhPewYpFl9X6HlWFBuASCTx2qgGizhpZhmn0,862
-llm_utils/vector_cache/cli.py,sha256=DMXTj8nZ2_LRjprbYPb4uzq04qZtOfBbmblmaqDcCuM,6251
-llm_utils/vector_cache/core.py,sha256=J8ocRX9sBfzboQkf5vFF2cx0SK-nftmKWJUa91WUBy8,31134
-llm_utils/vector_cache/types.py,sha256=ru8qmUZ8_lNd3_oYpjCMtpXTsqmwsSBe56Z4hTWm3xI,435
-llm_utils/vector_cache/utils.py,sha256=dwbbXlRrARrpmS4YqSlYQqrTURg0UWe8XvaAWcX05MM,1458
-speedy_utils/__init__.py,sha256=wPz1MNAicV7skqqZloUFt5QrJcAhxtPQ4jFXk2lz6YA,6190
-speedy_utils/all.py,sha256=gXXRlBLvU8AON7XqO6iFQ8LCIQEIcP_2CDumd_U1ppI,5171
-speedy_utils/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-speedy_utils/common/clock.py,sha256=3n4FkCW0dz46O8By09V5Pve1DSMgpLDRbWEVRryryeQ,7423
-speedy_utils/common/function_decorator.py,sha256=BspJ0YuGL6elS7lWBAgELZ-sCfED_1N2P5fgH-fCRUQ,2132
-speedy_utils/common/logger.py,sha256=a2iZx0eWyfi2-2X_H2QmfuA3tfR7_XSM7Nd0GdUnUOs,6435
-speedy_utils/common/notebook_utils.py,sha256=-97kehJ_Gg3TzDLubsLIYJcykqX1NXhbvBO6nniZSYM,2063
-speedy_utils/common/patcher.py,sha256=VCmdxyTF87qroggQkQklRPhAOPJbeBqhcJoTsLcDxNw,2303
-speedy_utils/common/report_manager.py,sha256=eBiw5KY6bWUhwki3B4lK5o8bFsp7L5x28X9GCI-Sd1w,3899
-speedy_utils/common/utils_cache.py,sha256=h3JbIi0V5pTaFNJDjfwORSN63bc0SrRq_dm8KZJiL94,27023
-speedy_utils/common/utils_io.py,sha256=E7mbxB_OpLvNWoFM2Qpxi1jaD8VwF-tvNOpGbf7swuU,14849
-speedy_utils/common/utils_misc.py,sha256=yYlyP0eXQuapY1dn5O8-UDePPq5bb6FxKFjb1kfZy5o,2354
-speedy_utils/common/utils_print.py,sha256=syRrnSFtguxrV-elx6DDVcSGu4Qy7D_xVNZhPwbUY4A,4864
-speedy_utils/multi_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-speedy_utils/multi_worker/process.py,sha256=RGGGnbZXCbEbdmxFVmnNfyccClAlflzRPE0d1C3CeeE,11385
-speedy_utils/multi_worker/thread.py,sha256=bRjxUHkBjbXHQ2KSsf-Zao28zbSId-8mqMFHwSG1l1s,25206
-speedy_utils/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-speedy_utils/scripts/mpython.py,sha256=IvywP7Y0_V6tWfMP-4MjPvN5_KfxWF21xaLJsCIayCk,3821
-speedy_utils/scripts/openapi_client_codegen.py,sha256=f2125S_q0PILgH5dyzoKRz7pIvNEjCkzpi4Q4pPFRZE,9683
-speedy_utils-1.1.26.dist-info/METADATA,sha256=j0mnIF3a8jD-WBaZgrXmIlWnzdhci8ncem57-oZt8cw,8028
-speedy_utils-1.1.26.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-speedy_utils-1.1.26.dist-info/entry_points.txt,sha256=1rrFMfqvaMUE9hvwGiD6vnVh98kmgy0TARBj-v0Lfhs,244
-speedy_utils-1.1.26.dist-info/RECORD,,