change it inlcude data and support os.Pathlike

Author: inikishev

pyproject.toml

@@ -26,4 +26,10 @@ dependencies = [
 
 [tool.setuptools.packages]
 find = {}  # Scanning implicit namespaces is active by default
-#1
+#1
+
+[tool.setuptools.package-data]
+# The key is the package name.
+# The value is a list of file patterns to include.
+# "data/**/*" means include all files in the 'data' subdirectory, recursively.
+visualbench = ["data/**/*"]

tests/__init__.py

@@ -1,5 +1,7 @@
 import os
 from collections.abc import Sequence
+from importlib import resources
+from pathlib import Path
 
 import numpy as np
 import torch
@@ -8,48 +10,51 @@
 from ..utils import normalize, to_3HW
 from ..utils.image import _imread
 
-_path = os.path.dirname(__file__)
+# _path = os.path.dirname(__file__)
 
-QRCODE96 = os.path.join(_path, 'qr-96.jpg')
+def _get_path(fname:str):
+    return Path(str(resources.files("visualbench").joinpath("data", fname)))
+
+QRCODE96 = _get_path('qr-96.jpg')
 """QR code that links to my github account"""
 
-ATTNGRAD96 = os.path.join(_path, 'attngrad-96.png')
+ATTNGRAD96 = _get_path('attngrad-96.png')
 """Piece of gradient of some model from transformers library except I don't remember which one"""
 
-SANIC96 = os.path.join(_path, 'sanic-96.jpg')
+SANIC96 = _get_path('sanic-96.jpg')
 """is an 8 year old image from my images folder and I think it is a screenshot from one of the sanic games"""
 
-FROG96 = os.path.join(_path, 'frog-96.png')
+FROG96 = _get_path('frog-96.png')
 """frame from https://www.youtube.com/@NinjaFrog777/videos"""
 
-WEEVIL96 = os.path.join(_path, 'weevil-96.png')
+WEEVIL96 = _get_path('weevil-96.png')
 """is from http://growingorganic.com/ipm-guide/weevils/"""
 
-TEST96 = os.path.join(_path, 'test-96.jpg')
+TEST96 = _get_path('test-96.jpg')
 """this was generated in like 2012 ago by google doodle generator and its still my favourite image and it is called test"""
 
-MAZE96 = os.path.join(_path, 'maze-96.png')
+MAZE96 = _get_path('maze-96.png')
 """a generic maze"""
 
-TEXT96 = os.path.join(_path, 'text-96.png')
+TEXT96 = _get_path('text-96.png')
 """lorem ipsum from lorem ipsum text"""
 
-GEOM96 = os.path.join(_path, 'geometry-96.png')
+GEOM96 = _get_path('geometry-96.png')
 """CC0 image from wikicommons, SORRY I CAN'T FIND THE LINK ANYMORE!!!"""
 
-RUBIC96 = os.path.join(_path, 'rubic-96.png')
+RUBIC96 = _get_path('rubic-96.png')
 """is from https://speedsolving.fandom.com/wiki/Rubik%27s_Cube?file=Rubik%27s_Cube_transparency.png"""
 
-SPIRAL96 = os.path.join(_path, 'spiral-96.png')
+SPIRAL96 = _get_path('spiral-96.png')
 """A colorful spiral"""
 
-BIANG96 = os.path.join(_path, 'biang-96.png')
+BIANG96 = _get_path('biang-96.png')
 """apparently its the hardest hieroglyph and it is from https://commons.wikimedia.org/wiki/File:Bi%C3%A1ng_%28regular_script%29.svg"""
 
-EMOJIS96 = os.path.join(_path, 'emojis-96.png')
+EMOJIS96 = _get_path('emojis-96.png')
 """some random emojis"""
 
-GRID96 = os.path.join(_path, 'grid-96.png')
+GRID96 = _get_path('grid-96.png')
 """Grid of black and white cells"""
 
 def get_qrcode():
@@ -165,7 +170,7 @@ def get_ill_conditioned(size: int | tuple[int,int], cond:float=1e17):
 
 def get_font_dict(dtype=torch.bool, device=None):
     """returns a dictionary which maps letters, numbers and +-*/|. to 3x3 binary images."""
-    path = os.path.join(_path, '3x3 font.jpeg')
+    path = _get_path('3x3 font.jpeg')
     image = to_3HW(_imread(path).float()).mean(0)
     image = torch.where(image > 128, 1, 0).contiguous().to(dtype=dtype, device=device)
 

visualbench/data/__init__.py

@@ -1,11 +1,13 @@
 import warnings
+from os import PathLike
 from typing import Any
 
 import numpy as np
 import torch
 
 from .image import _imread
 
+
 def _imread_normalize(x) -> torch.Tensor:
     x = _imread(x).float()
     x -= x.mean()
@@ -22,7 +24,7 @@ def is_scalar(x: Any) -> bool:
     return isinstance(x, (int,float,bool))
 
 def totensor(x, device=None, dtype=None, clone=None) -> torch.Tensor:
-    if isinstance(x, str): x = _imread_normalize(x)
+    if isinstance(x, (str, PathLike)): x = _imread_normalize(x)
 
     if isinstance(x, torch.Tensor): x = x.to(dtype=dtype, device=device, copy=False)
 
@@ -42,7 +44,7 @@ def totensor(x, device=None, dtype=None, clone=None) -> torch.Tensor:
     return x
 
 def tonumpy(x) -> np.ndarray:
-    if isinstance(x, str): x = _imread_normalize(x)
+    if isinstance(x, (str, PathLike)): x = _imread_normalize(x)
     if isinstance(x, np.ndarray): return x
     if isinstance(x, torch.Tensor): return x.numpy(force=True)
     return np.asarray(x)

visualbench/utils/format.py

@@ -1,35 +1,39 @@
 import warnings
+from os import PathLike
+
 import numpy as np
 import torch
 
-def _imread_skimage(path:str) -> np.ndarray:
+
+def _imread_skimage(path:str|PathLike) -> np.ndarray:
     import skimage
     return skimage.io.imread(path)
 
-def _imread_plt(path:str) -> np.ndarray:
+def _imread_plt(path:str|PathLike) -> np.ndarray:
     import matplotlib.pyplot as plt
-    return plt.imread(path)
+    return plt.imread(str(path))
 
-def _imread_cv2(path):
+def _imread_cv2(path: str|PathLike) -> np.ndarray:
     import cv2
-    image = cv2.imread(path) # pylint:disable=no-member
+    image = cv2.imread(str(path)) # pylint:disable=no-member
+    assert image is not None
     if image.ndim == 3: image = image[:, :, ::-1] # BRG -> RGB
     return image
 
-def _imread_imageio(path):
+def _imread_imageio(path: str|PathLike):
     from imageio import v3
-    return v3.imread(path)
+    return v3.imread(str(path))
 
-def _imread_pil(path:str) -> np.ndarray:
+def _imread_pil(path: str|PathLike) -> np.ndarray:
     import PIL.Image
     return np.array(PIL.Image.open(path))
 
-def _imread_torchvision(path:str, dtype=None, device=None) -> torch.Tensor:
+def _imread_torchvision(path: str|PathLike, dtype=None, device=None) -> torch.Tensor:
     import torchvision
-    return torchvision.io.read_image(path).to(dtype=dtype, device=device, copy=
+    return torchvision.io.read_image(str(path)).to(dtype=dtype, device=device, copy=
                                               False)
 
-def _imread(path: str) -> torch.Tensor:
+def _imread(path: str|PathLike) -> torch.Tensor:
     try: return _imread_torchvision(path)
     except Exception:
         img = None