<filename>src/utils/__init__.py
import re
import torch
from torch import Tensor as TorchTensor
from pathlib import Path
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
from .TensorboardLogger import TensorboardLogger  # Keep this even if unused, it exposes Tensorboard Logger in utils
import cv2
from urllib.request import urlretrieve
from zipfile import ZipFile
from tqdm import tqdm
import ray
import os
from typing import Optional

epsilon_f32 = np.finfo(np.float32).eps

def indices_to_mask(indices, size):
    """Converts a tensor of indices into a bool mask 1D tensor of the given size"""
    mask = torch.zeros((size,), dtype=torch.bool).to(indices.device)
    mask.scatter_(0, indices, True)
    return mask

"""dtype conversions"""
def numpy_image_to_torch(x: np.ndarray):
    """Converts a (B)HWC uint8 RGB image from numpy to torch (B)CHW float format."""
    assert isinstance(x, np.ndarray), f"Illegal input type '{type(x)}'"
    if x.ndim == 4 and x.shape[1] != 3:
        x = x.transpose((0,3,1,2)) # BHWC -> BCHW
    elif x.ndim == 3 and x.shape[0] != 3:
        x = x.transpose((2,0,1)) # HWC -> CHW
    if x.dtype == np.uint8:
        x = x.astype(np.float32) / 255.0
    x = torch.from_numpy(x)
    return x

def torch_image_to_numpy(x: TorchTensor):
    """Converts a (B)CHW float RGB image from torch to numpy format (B)HWC (dtype stays float)."""
    assert isinstance(x, TorchTensor), f"Illegal input type '{type(x)}'"
    if x.ndim == 4 and x.shape[-1] != 3:
        x = x.permute(0,2,3,1) # BHWC -> BCHW
    elif x.ndim == 3 and x.shape[-1] != 3:
        x = x.permute(1,2,0) # HWC -> CHW
    if x.dtype == torch.uint8:
        x = x / 255.0
    x = x.cpu().numpy()
    return x

"""Image reading and plotting utilities"""
def imread(p: Path, rgb: bool = True):
    if type(p) == str:
        p = Path(p)
    if p.suffix.endswith(".npy"):
        x = np.load(p)
    else:
        x = Image.open(p)
        if rgb:
            x = x.convert('RGB')
        x = np.array(x)
    if x.dtype == np.uint8 or x.dtype == np.uint16:
        x = x.astype(np.float32) / 255.0
    return x

def imshow(img, title=None, axis_off=True, new_figure=True, **kwargs):
    if isinstance(img, Path) or type(img) == str:
        img = imread(img)
    if isinstance(img, TorchTensor):
        img = torch_image_to_numpy(img.cpu())
    if img.ndim == 3 and img.shape[0] == 3:
        img = img.transpose((1,2,0)) # CHW --> HWC
    f = plt.figure() if new_figure else None
    if axis_off:
        plt.axis("off")
    if title is not None:
        plt.title(title)
    plt.imshow(img, interpolation='none', **kwargs)
    return f

def imsave(path, img):
    if isinstance(img, Path) or type(img) == str:
        img = imread(img)
    if isinstance(img, TorchTensor):
        img = torch_image_to_numpy(img.cpu())
    if img.ndim == 3 and img.shape[0] == 3:
        img = img.transpose((1,2,0)) # CHW --> HWC
    plt.imsave(path, img)


def apply_color_mask(rgb, mask, color=np.array([0,1.0,0]), alpha=0.3):
    assert rgb.ndim == 3, f"Input shape is {rgb}. Only 3 dimensions are supported"
    was_torch =  isinstance(rgb, TorchTensor)
    if was_torch:
        rgb = torch_image_to_numpy(rgb)
    if rgb.ndim == 2:
        rgb = np.stack([rgb,rgb,rgb], axis=-1)
    mask = np.stack([mask, mask, mask], axis=-1)
    h,w = rgb.shape[:2]
    uniform_color = np.tile(color[None,None,:], [h,w,1])
    if uniform_color.dtype == np.uint8 and uniform_color.dtype != np.uint8:
        uniform_color = uniform_color / 255.
    out = rgb.copy()
    out[mask] = out[mask]*(1-alpha) + uniform_color[mask]*alpha
    if was_torch:
        out = numpy_image_to_torch(out)
    return out

def show_pred(index, rgb, pred, gt, mask, **kwargs):
    rgb = apply_color_mask(rgb, pred & gt & mask, np.array((0,1.,0))) # good = green
    rgb = apply_color_mask(rgb, pred & ~gt & mask, np.array((1.,0,0))) # bad = red
    rgb = apply_color_mask(rgb, pred & ~mask, np.array((0.0,0.0,1.0))) # Good out of mask = blue
    imshow(rgb, **kwargs)


"""Extensions to pathlib.Path"""
def lglob(self: Path, pattern="*"):
    """Like Path.glob, but returns a list rather than a generator"""
    return list(self.glob(pattern))

def rglob(self: Path, regex=".*"):
    """Like path.glob, but uses a regex to match Paths"""
    return (f for f in self.glob("*") if re.match(regex, str(f)))

def lrglob(self: Path, pattern=".*"):
    """Like path.rglob, but returns a list rather than a generator"""
    return list(self.rglob(pattern))

Path.lglob = lglob
Path.rglob = rglob
Path.lrglob = lrglob


"""File utilities"""
def download_file_if_not_exists(path: Path, url: str):
    if not path.exists():
        path.parent.mkdir(exist_ok=True, parents=True)
        urlretrieve(url, str(path))
    return path
        
def unzip(path: Path, extract_dir: Optional[Path] = None, delete_archive: bool = False):
    if extract_dir is None: extract_dir = path.parent / path.name
    extract_dir.mkdir(exist_ok=True)
    with ZipFile(path, 'r') as zf:
        zf.extractall(extract_dir)
    if delete_archive:
        path.unlink()


"""Parallel CPU Processing"""
def tqdm_ray(it, *args, **kwargs):
    def to_iterator(obj_ids):
        while obj_ids:
            done, obj_ids = ray.wait(obj_ids)
            yield ray.get(done[0])
    return tqdm(to_iterator(it), *args, **kwargs)
    

def key_to_tuple(k):
    t = tuple(k.split(","))
    return t if len(t) > 1 else t[0]

def tuple_to_key(t):
    return "_".join(t)

def assert_writable(p: Path):
    assert os.access(p, os.W_OK), f"Can't write to directory '{p}'"

<filename>src/utils/__init__.py
import re
import torch
from torch import Tensor as TorchTensor
from pathlib import Path
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
from .TensorboardLogger import TensorboardLogger  # Keep this even if unused, it exposes Tensorboard Logger in utils
import cv2
from urllib.request import urlretrieve
from zipfile import ZipFile
from tqdm import tqdm
import ray
import os
from typing import Optional

epsilon_f32 = np.finfo(np.float32).eps

def indices_to_mask(indices, size):
    """Converts a tensor of indices into a bool mask 1D tensor of the given size"""
    mask = torch.zeros((size,), dtype=torch.bool).to(indices.device)
    mask.scatter_(0, indices, True)
    return mask

"""dtype conversions"""
def numpy_image_to_torch(x: np.ndarray):
    """Converts a (B)HWC uint8 RGB image from numpy to torch (B)CHW float format."""
    assert isinstance(x, np.ndarray), f"Illegal input type '{type(x)}'"
    if x.ndim == 4 and x.shape[1] != 3:
        x = x.transpose((0,3,1,2)) # BHWC -> BCHW
    elif x.ndim == 3 and x.shape[0] != 3:
        x = x.transpose((2,0,1)) # HWC -> CHW
    if x.dtype == np.uint8:
        x = x.astype(np.float32) / 255.0
    x = torch.from_numpy(x)
    return x

def torch_image_to_numpy(x: TorchTensor):
    """Converts a (B)CHW float RGB image from torch to numpy format (B)HWC (dtype stays float)."""
    assert isinstance(x, TorchTensor), f"Illegal input type '{type(x)}'"
    if x.ndim == 4 and x.shape[-1] != 3:
        x = x.permute(0,2,3,1) # BHWC -> BCHW
    elif x.ndim == 3 and x.shape[-1] != 3:
        x = x.permute(1,2,0) # HWC -> CHW
    if x.dtype == torch.uint8:
        x = x / 255.0
    x = x.cpu().numpy()
    return x

"""Image reading and plotting utilities"""
def imread(p: Path, rgb: bool = True):
    if type(p) == str:
        p = Path(p)
    if p.suffix.endswith(".npy"):
        x = np.load(p)
    else:
        x = Image.open(p)
        if rgb:
            x = x.convert('RGB')
        x = np.array(x)
    if x.dtype == np.uint8 or x.dtype == np.uint16:
        x = x.astype(np.float32) / 255.0
    return x

def imshow(img, title=None, axis_off=True, new_figure=True, **kwargs):
    if isinstance(img, Path) or type(img) == str:
        img = imread(img)
    if isinstance(img, TorchTensor):
        img = torch_image_to_numpy(img.cpu())
    if img.ndim == 3 and img.shape[0] == 3:
        img = img.transpose((1,2,0)) # CHW --> HWC
    f = plt.figure() if new_figure else None
    if axis_off:
        plt.axis("off")
    if title is not None:
        plt.title(title)
    plt.imshow(img, interpolation='none', **kwargs)
    return f

def imsave(path, img):
    if isinstance(img, Path) or type(img) == str:
        img = imread(img)
    if isinstance(img, TorchTensor):
        img = torch_image_to_numpy(img.cpu())
    if img.ndim == 3 and img.shape[0] == 3:
        img = img.transpose((1,2,0)) # CHW --> HWC
    plt.imsave(path, img)


def apply_color_mask(rgb, mask, color=np.array([0,1.0,0]), alpha=0.3):
    assert rgb.ndim == 3, f"Input shape is {rgb}. Only 3 dimensions are supported"
    was_torch =  isinstance(rgb, TorchTensor)
    if was_torch:
        rgb = torch_image_to_numpy(rgb)
    if rgb.ndim == 2:
        rgb = np.stack([rgb,rgb,rgb], axis=-1)
    mask = np.stack([mask, mask, mask], axis=-1)
    h,w = rgb.shape[:2]
    uniform_color = np.tile(color[None,None,:], [h,w,1])
    if uniform_color.dtype == np.uint8 and uniform_color.dtype != np.uint8:
        uniform_color = uniform_color / 255.
    out = rgb.copy()
    out[mask] = out[mask]*(1-alpha) + uniform_color[mask]*alpha
    if was_torch:
        out = numpy_image_to_torch(out)
    return out

def show_pred(index, rgb, pred, gt, mask, **kwargs):
    rgb = apply_color_mask(rgb, pred & gt & mask, np.array((0,1.,0))) # good = green
    rgb = apply_color_mask(rgb, pred & ~gt & mask, np.array((1.,0,0))) # bad = red
    rgb = apply_color_mask(rgb, pred & ~mask, np.array((0.0,0.0,1.0))) # Good out of mask = blue
    imshow(rgb, **kwargs)


"""Extensions to pathlib.Path"""
def lglob(self: Path, pattern="*"):
    """Like Path.glob, but returns a list rather than a generator"""
    return list(self.glob(pattern))

def rglob(self: Path, regex=".*"):
    """Like path.glob, but uses a regex to match Paths"""
    return (f for f in self.glob("*") if re.match(regex, str(f)))

def lrglob(self: Path, pattern=".*"):
    """Like path.rglob, but returns a list rather than a generator"""
    return list(self.rglob(pattern))

Path.lglob = lglob
Path.rglob = rglob
Path.lrglob = lrglob


"""File utilities"""
def download_file_if_not_exists(path: Path, url: str):
    if not path.exists():
        path.parent.mkdir(exist_ok=True, parents=True)
        urlretrieve(url, str(path))
    return path
        
def unzip(path: Path, extract_dir: Optional[Path] = None, delete_archive: bool = False):
    if extract_dir is None: extract_dir = path.parent / path.name
    extract_dir.mkdir(exist_ok=True)
    with ZipFile(path, 'r') as zf:
        zf.extractall(extract_dir)
    if delete_archive:
        path.unlink()


"""Parallel CPU Processing"""
def tqdm_ray(it, *args, **kwargs):
    def to_iterator(obj_ids):
        while obj_ids:
            done, obj_ids = ray.wait(obj_ids)
            yield ray.get(done[0])
    return tqdm(to_iterator(it), *args, **kwargs)
    

def key_to_tuple(k):
    t = tuple(k.split(","))
    return t if len(t) > 1 else t[0]

def tuple_to_key(t):
    return "_".join(t)

def assert_writable(p: Path):
    assert os.access(p, os.W_OK), f"Can't write to directory '{p}'"

# Keep this even if unused, it exposes Tensorboard Logger in utils Converts a tensor of indices into a bool mask 1D tensor of the given size dtype conversions Converts a (B)HWC uint8 RGB image from numpy to torch (B)CHW float format. # BHWC -> BCHW # HWC -> CHW Converts a (B)CHW float RGB image from torch to numpy format (B)HWC (dtype stays float). # BHWC -> BCHW # HWC -> CHW Image reading and plotting utilities # CHW --> HWC # CHW --> HWC # good = green # bad = red # Good out of mask = blue Extensions to pathlib.Path Like Path.glob, but returns a list rather than a generator Like path.glob, but uses a regex to match Paths Like path.rglob, but returns a list rather than a generator File utilities Parallel CPU Processing