opentau 0.1.0__py3-none-any.whl

opentau/datasets/grounding/pixmo.py

@@ -0,0 +1,177 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Datasets for Image-Text Point Set grounding tasks.
+
+This module provides the PIXMO (Pixel-level Manipulation) dataset implementation
+for training vision-language models on part localization and object grounding tasks.
+"""
+
+import json
+import logging
+import random
+import warnings
+from io import BytesIO
+
+import numpy as np
+import requests
+import torch
+from datasets import load_dataset
+from PIL import Image
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from opentau import register_grounding_dataset
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.grounding.base import GroundingDataset
+
+# TODO: add a config to filter the warnings
+logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
+warnings.filterwarnings(
+    "ignore",
+    message=r"Palette images with Transparency expressed in bytes should be converted to RGBA images",
+    category=UserWarning,
+    module=r"PIL\.Image",
+)
+warnings.filterwarnings(
+    "ignore",
+    message=r"image file could not be identified because AVIF support not installed",
+    category=UserWarning,
+    module=r"PIL\.Image",
+)
+
+IMG_SIZE = 224
+POINT_GRID = 255
+MAX_RETRIES = 1
+HTTP_TIMEOUT = 1
+LOG_EVERY_N_BAD = 1000
+
+_session = requests.Session()
+_session.mount(
+    "https://",
+    HTTPAdapter(
+        max_retries=Retry(
+            total=MAX_RETRIES,
+            backoff_factor=0.5,
+            status_forcelist=[500, 502, 503, 504],
+        )
+    ),
+)
+
+
+def _pil_from_url(url: str) -> Image.Image | None:
+    """Download, decode, and resize an image using its URL. Returns None in case of failure."""
+    try:
+        r = _session.get(url, timeout=HTTP_TIMEOUT)
+        r.raise_for_status()
+        # TODO: Check against the hash in case the image somehow changed.
+        return Image.open(BytesIO(r.content)).convert("RGB")
+    except Exception:
+        return None
+
+
+def _get_post_fix(label: str, points: list, orig_w: int, orig_h: int, max_points: int = 16) -> str:
+    """Map points from pixel space to grid space and return a JSON postfix string.
+
+    Converts pixel coordinates to a 255x255 grid, deduplicates points, and
+    limits to max_points. Returns a JSON string with point coordinates and labels.
+
+    Args:
+        label: Label for the points (e.g., object class name).
+        points: List of point dictionaries with 'x' and 'y' keys.
+        orig_w: Original image width.
+        orig_h: Original image height.
+        max_points: Maximum number of points to include. Defaults to 16.
+
+    Returns:
+        JSON string containing point coordinates and labels.
+    """
+    # use `dict` to deduplicate as `set` is not guaranteed to preserve order
+    deduplicated = {
+        (int(p["x"] * POINT_GRID / orig_w), int(p["y"] * POINT_GRID / orig_h)): None for p in points
+    }
+    if len(deduplicated) > max_points:
+        deduplicated = random.choices(list(deduplicated), k=max_points)
+    rows = [{"in_frame": True, "point": pair, "label": label} for pair in deduplicated]
+    return json.dumps(rows)
+
+
+def _img_to_normalized_tensor(img: Image.Image) -> torch.Tensor:
+    """Convert a PIL Image to a normalized torch tensor.
+
+    Resizes the image to IMG_SIZE and converts it from (H, W, C) to (C, H, W)
+    format, normalizing pixel values to [0, 1].
+
+    Args:
+        img: PIL Image to convert.
+
+    Returns:
+        Normalized tensor of shape (C, IMG_SIZE, IMG_SIZE) with values in [0, 1].
+    """
+    img = img.resize((IMG_SIZE, IMG_SIZE), Image.BILINEAR)
+    # pytorch uses (C, H, W) while PIL uses (H, W, C)
+    return torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0
+
+
+@register_grounding_dataset("pixmo")
+class PixmoDataset(GroundingDataset):
+    r"""Dataset for the iterable PixMo dataset implementation, recommended to be used together with PrefetchWrapper"""
+
+    def __init__(self, cfg: TrainPipelineConfig, consecutive_bad_tolerance=100):
+        # Self.ds is needed for metadata, which is computed in parent constructor
+        self.ds = load_dataset("allenai/pixmo-points", split="train")
+        super().__init__(cfg)
+        self.bad_ids = set()
+        self.consecutive_bad_tolerance = consecutive_bad_tolerance
+
+    def __len__(self):
+        return len(self.ds)
+
+    def _get_feature_mapping_key(self) -> str:
+        return "pixmo"
+
+    def __getitem_helper__(self, item) -> dict:
+        """Get a PixMo dataset item.
+
+        Downloads the image from URL and formats it for part localization tasks.
+        Retries with random indices if image download fails.
+
+        Args:
+            item: Index of the item to retrieve.
+
+        Returns:
+            Dictionary with image, task, postfix, task_type, and prompt.
+
+        Raises:
+            RuntimeError: If too many consecutive items fail to load.
+        """
+        for _ in range(self.consecutive_bad_tolerance):
+            if item in self.bad_ids:
+                item = np.random.randint(0, len(self.ds))
+                continue
+            ex = self.ds[item]
+            img = _pil_from_url(ex["image_url"])
+            if img is None:
+                self.bad_ids.add(item)
+                item = np.random.randint(0, len(self.ds))
+                continue
+
+            return {
+                "image": _img_to_normalized_tensor(img),
+                "task": ex["label"],
+                "postfix": _get_post_fix(ex["label"], ex["points"], *img.size),
+                "task_type": "part",
+                "prompt": f'{{"task": "part", "description": "Find {ex["label"]} in the image"}}',
+            }
+
+        raise RuntimeError("Too many consecutive bad items. Please check dataset or increase the tolerance.")

opentau/datasets/grounding/vsr.py

@@ -0,0 +1,141 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""VSR (Visual Spatial Reasoning) dataset for true/false statement grounding.
+
+This module provides the VSR dataset implementation for training vision-language
+models on visual spatial reasoning tasks. The dataset contains images with
+statements about spatial relationships, and models must determine whether each
+statement is true or false based on the image content.
+"""
+
+import logging
+from io import BytesIO
+
+import numpy as np
+import requests
+import torch
+from datasets import load_dataset
+from PIL import Image
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from opentau import register_grounding_dataset
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.grounding.base import GroundingDataset
+
+logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
+
+MAX_RETRIES = 1
+HTTP_TIMEOUT = 1
+LOG_EVERY_N_BAD = 1000
+
+_session = requests.Session()
+_session.mount(
+    "https://",
+    HTTPAdapter(
+        max_retries=Retry(
+            total=MAX_RETRIES,
+            backoff_factor=0.5,
+            status_forcelist=[500, 502, 503, 504],
+        )
+    ),
+)
+
+
+def _pil_from_url(url: str) -> Image.Image | None:
+    """Download, decode, and resize an image using its URL. Returns None in case of failure."""
+    try:
+        r = _session.get(url, timeout=HTTP_TIMEOUT)
+        r.raise_for_status()
+        # TODO: Check against the hash in case the image somehow changed.
+        return Image.open(BytesIO(r.content)).convert("RGB")
+    except Exception:
+        return None
+
+
+def _img_to_normalized_tensor(img: Image.Image, img_shape: tuple) -> torch.Tensor:
+    """Convert a PIL Image to a normalized torch tensor.
+
+    Resizes the image and converts it from (H, W, C) to (C, H, W) format,
+    normalizing pixel values to [0, 1].
+
+    Args:
+        img: PIL Image to convert.
+        img_shape: Target image shape (height, width).
+
+    Returns:
+        Normalized tensor of shape (C, H, W) with values in [0, 1].
+    """
+    img = img.resize(img_shape, Image.BILINEAR)
+    return torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0
+
+
+@register_grounding_dataset("vsr")
+class VSRDataset(GroundingDataset):
+    """Visual Spatial Reasoning (VSR) dataset for true/false statement grounding.
+
+    Loads the cambridgeltl/vsr_random dataset from HuggingFace and formats it
+    for visual reasoning tasks where models must determine if statements about
+    images are true or false.
+    """
+
+    def __init__(self, cfg: TrainPipelineConfig, consecutive_bad_tolerance=100):
+        self.dataset = load_dataset("cambridgeltl/vsr_random", split="train")
+        super().__init__(cfg)
+        self.bad_ids = set()
+        self.consecutive_bad_tolerance = consecutive_bad_tolerance
+        self.mapping = {0: "False", 1: "True"}
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def _get_feature_mapping_key(self) -> str:
+        return "vsr"
+
+    def __getitem_helper__(self, item) -> dict:
+        """Get a VSR dataset item.
+
+        Downloads the image from URL and formats it for true/false reasoning tasks.
+        Retries with random indices if image download fails.
+
+        Args:
+            item: Index of the item to retrieve.
+
+        Returns:
+            Dictionary with image, task, postfix, task_type, and prompt.
+
+        Raises:
+            RuntimeError: If too many consecutive items fail to load.
+        """
+        for _ in range(self.consecutive_bad_tolerance):
+            if item in self.bad_ids:
+                item = np.random.randint(0, len(self.dataset))
+                continue
+
+            sample = self.dataset[item]
+            img = _pil_from_url(sample["image_link"])
+            if img is None:
+                self.bad_ids.add(item)
+                item = np.random.randint(0, len(self.dataset))
+                continue
+
+            return {
+                "image": _img_to_normalized_tensor(img, self.resolution),
+                "task": sample["label"],
+                "postfix": f"The statement is {self.mapping[sample['label']]}",
+                "task_type": "grounding",
+                "prompt": f'{{"task": "grounding", "description": "Using the Image, Tell me if following statement is true or false. \n  {sample["caption"]}"}}',
+            }
+
+        raise RuntimeError("Too many consecutive bad items. Please check dataset or increase the tolerance.")

opentau/datasets/image_writer.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Asynchronous image writing utilities for high-frequency data recording.
+
+This module provides functionality for writing images to disk asynchronously
+using multithreading or multiprocessing, which is critical for controlling
+robots and recording data at high frame rates without blocking the main process.
+
+The module supports two execution models:
+
+    1. Threading mode (num_processes=0): Creates a pool of worker threads
+       for concurrent image writing within a single process.
+    2. Multiprocessing mode (num_processes>0): Creates multiple processes,
+       each with their own thread pool, for maximum parallelism.
+
+Key Features:
+    - Asynchronous writing: Images are queued and written in background
+      workers, preventing I/O blocking of the main process.
+    - Multiple input formats: Supports torch Tensors, numpy arrays, and
+      PIL Images with automatic conversion.
+    - Format flexibility: Handles both channel-first (C, H, W) and
+      channel-last (H, W, C) image formats.
+    - Type conversion: Automatically converts float arrays in [0, 1] to
+      uint8 in [0, 255] for PIL Image compatibility.
+    - Safe cleanup: Decorator ensures image writers are properly stopped
+      even when exceptions occur.
+
+Classes:
+
+    AsyncImageWriter
+        Main class for asynchronous image writing with configurable threading
+        or multiprocessing backends.
+
+Functions:
+    image_array_to_pil_image
+        Convert numpy array to PIL Image with format and type conversion.
+    write_image
+        Write an image (numpy array or PIL Image) to disk.
+    worker_thread_loop
+        Worker thread loop for processing image write queue.
+    worker_process
+        Worker process that manages multiple threads for image writing.
+    safe_stop_image_writer
+        Decorator to safely stop image writer on exceptions.
+
+Example:
+    Create an async image writer with threading:
+        >>> writer = AsyncImageWriter(num_processes=0, num_threads=4)
+        >>> writer.save_image(image_array, Path("output/image.jpg"))
+        >>> writer.wait_until_done()  # Wait for all images to be written
+        >>> writer.stop()  # Clean up resources
+"""
+
+import multiprocessing
+import queue
+import threading
+from pathlib import Path
+
+import numpy as np
+import PIL.Image
+import torch
+
+
+def safe_stop_image_writer(func):
+    """Decorator to safely stop image writer on exceptions.
+
+    Ensures that the image writer is properly stopped if an exception occurs
+    during function execution.
+
+    Args:
+        func: Function to wrap.
+
+    Returns:
+        Wrapped function that stops image writer on exceptions.
+    """
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            dataset = kwargs.get("dataset")
+            image_writer = getattr(dataset, "image_writer", None) if dataset else None
+            if image_writer is not None:
+                print("Waiting for image writer to terminate...")
+                image_writer.stop()
+            raise e
+
+    return wrapper
+
+
+def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) -> PIL.Image.Image:
+    """Convert a numpy array to a PIL Image.
+
+    Supports channel-first (C, H, W) and channel-last (H, W, C) formats.
+    Converts float arrays in [0, 1] to uint8 in [0, 255].
+
+    Args:
+        image_array: Input image array of shape (C, H, W) or (H, W, C).
+        range_check: If True, validates that float arrays are in [0, 1] range.
+            Defaults to True.
+
+    Returns:
+        PIL Image object.
+
+    Raises:
+        ValueError: If array has wrong number of dimensions, wrong number of
+            channels, or float values are outside [0, 1] range.
+        NotImplementedError: If image doesn't have 3 channels.
+
+    Note:
+        TODO(aliberts): handle 1 channel and 4 for depth images
+    """
+    # TODO(aliberts): handle 1 channel and 4 for depth images
+    if image_array.ndim != 3:
+        raise ValueError(f"The array has {image_array.ndim} dimensions, but 3 is expected for an image.")
+
+    if image_array.shape[0] == 3:
+        # Transpose from pytorch convention (C, H, W) to (H, W, C)
+        image_array = image_array.transpose(1, 2, 0)
+
+    elif image_array.shape[-1] != 3:
+        raise NotImplementedError(
+            f"The image has {image_array.shape[-1]} channels, but 3 is required for now."
+        )
+
+    if image_array.dtype != np.uint8:
+        if range_check:
+            max_ = image_array.max().item()
+            min_ = image_array.min().item()
+            if max_ > 1.0 or min_ < 0.0:
+                raise ValueError(
+                    "The image data type is float, which requires values in the range [0.0, 1.0]. "
+                    f"However, the provided range is [{min_}, {max_}]. Please adjust the range or "
+                    "provide a uint8 image with values in the range [0, 255]."
+                )
+
+        image_array = (image_array * 255).astype(np.uint8)
+
+    return PIL.Image.fromarray(image_array)
+
+
+def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path) -> None:
+    """Write an image to disk.
+
+    Converts numpy arrays to PIL Images if needed, then saves to the specified path.
+
+    Args:
+        image: Image to save (numpy array or PIL Image).
+        fpath: Path where the image will be saved.
+
+    Raises:
+        TypeError: If image type is not supported.
+    """
+    try:
+        if isinstance(image, np.ndarray):
+            img = image_array_to_pil_image(image)
+        elif isinstance(image, PIL.Image.Image):
+            img = image
+        else:
+            raise TypeError(f"Unsupported image type: {type(image)}")
+        img.save(fpath)
+    except Exception as e:
+        print(f"Error writing image {fpath}: {e}")
+
+
+def worker_thread_loop(queue: queue.Queue) -> None:
+    """Worker thread loop for asynchronous image writing.
+
+    Continuously processes items from the queue until receiving None (sentinel).
+    Each item should be a tuple of (image_array, file_path).
+
+    Args:
+        queue: Queue containing (image_array, file_path) tuples or None sentinel.
+    """
+    while True:
+        item = queue.get()
+        if item is None:
+            queue.task_done()
+            break
+        image_array, fpath = item
+        write_image(image_array, fpath)
+        queue.task_done()
+
+
+def worker_process(queue: queue.Queue, num_threads: int) -> None:
+    """Worker process that manages multiple threads for image writing.
+
+    Creates and manages a pool of worker threads that process items from the queue.
+
+    Args:
+        queue: Queue containing (image_array, file_path) tuples or None sentinels.
+        num_threads: Number of worker threads to create in this process.
+    """
+    threads = []
+    for _ in range(num_threads):
+        t = threading.Thread(target=worker_thread_loop, args=(queue,))
+        t.daemon = True
+        t.start()
+        threads.append(t)
+    for t in threads:
+        t.join()
+
+
+class AsyncImageWriter:
+    """
+    This class abstract away the initialisation of processes or/and threads to
+    save images on disk asynchrounously, which is critical to control a robot and record data
+    at a high frame rate.
+
+    When `num_processes=0`, it creates a threads pool of size `num_threads`.
+    When `num_processes>0`, it creates processes pool of size `num_processes`, where each subprocess starts
+    their own threads pool of size `num_threads`.
+
+    The optimal number of processes and threads depends on your computer capabilities.
+    We advise to use 4 threads per camera with 0 processes. If the fps is not stable, try to increase or lower
+    the number of threads. If it is still not stable, try to use 1 subprocess, or more.
+    """
+
+    def __init__(self, num_processes: int = 0, num_threads: int = 1):
+        self.num_processes = num_processes
+        self.num_threads = num_threads
+        self.queue = None
+        self.threads = []
+        self.processes = []
+        self._stopped = False
+
+        if num_threads <= 0 and num_processes <= 0:
+            raise ValueError("Number of threads and processes must be greater than zero.")
+
+        if self.num_processes == 0:
+            # Use threading
+            self.queue = queue.Queue()
+            for _ in range(self.num_threads):
+                t = threading.Thread(target=worker_thread_loop, args=(self.queue,))
+                t.daemon = True
+                t.start()
+                self.threads.append(t)
+        else:
+            # Use multiprocessing
+            self.queue = multiprocessing.JoinableQueue()
+            for _ in range(self.num_processes):
+                p = multiprocessing.Process(target=worker_process, args=(self.queue, self.num_threads))
+                p.daemon = True
+                p.start()
+                self.processes.append(p)
+
+    def save_image(self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path) -> None:
+        """Queue an image for asynchronous writing.
+
+        Converts torch tensors to numpy arrays and adds the image to the write queue.
+
+        Args:
+            image: Image to save (torch Tensor, numpy array, or PIL Image).
+            fpath: Path where the image will be saved.
+        """
+        if isinstance(image, torch.Tensor):
+            # Convert tensor to numpy array to minimize main process time
+            image = image.cpu().numpy()
+        self.queue.put((image, fpath))
+
+    def wait_until_done(self) -> None:
+        """Wait until all queued images have been written to disk."""
+        self.queue.join()
+
+    def stop(self) -> None:
+        """Stop all worker threads/processes and clean up resources.
+
+        Sends sentinel values to all workers and waits for them to finish.
+        Terminates processes if they don't respond.
+        """
+        if self._stopped:
+            return
+
+        if self.num_processes == 0:
+            for _ in self.threads:
+                self.queue.put(None)
+            for t in self.threads:
+                t.join()
+        else:
+            num_nones = self.num_processes * self.num_threads
+            for _ in range(num_nones):
+                self.queue.put(None)
+            for p in self.processes:
+                p.join()
+                if p.is_alive():
+                    p.terminate()
+            self.queue.close()
+            self.queue.join_thread()
+
+        self._stopped = True