Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,087 Bytes
c8b42eb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 |
import os
import random
from functools import lru_cache
from typing import Tuple
import numpy as np
import pytest
import requests
import torch
from PIL import Image
from uniception.models.encoders import *
from uniception.models.encoders.image_normalizations import *
@pytest.fixture(scope="module")
def norm_types():
return IMAGE_NORMALIZATION_DICT.keys()
@pytest.fixture(scope="module")
def encoders():
return [
"croco",
"dust3r_224",
"dust3r_512",
"dust3r_512_dpt",
"mast3r_512",
"dinov2_base",
"dinov2_large",
"dinov2_large_reg",
"dinov2_large_dav2",
"dinov2_giant",
"dinov2_giant_reg",
"radio_v2.5-b",
"radio_v2.5-l",
"e-radio_v2",
"naradio_v2.5-b",
"naradio_v2.5-l",
"cosmosx8",
"patch_embedder",
]
@pytest.fixture(scope="module")
def encoder_configs(encoders):
# Adjust the number of configs to match the number of encoders
return [{}] * len(encoders)
@pytest.fixture
def device(request):
# Access the value of the custom option for device
device_str = request.config.getoption("--device")
if device_str == "gpu" and torch.cuda.is_available():
device = torch.device("cuda") # Use the default CUDA device
else:
device = torch.device("cpu")
print(f"Using device: {device.type.upper()}")
return device
@pytest.fixture
def example_input(device):
@lru_cache(maxsize=3)
def _get_example_input(
image_size: Tuple[int, int],
image_norm_type: str = "dummy",
img_selection: int = 1,
return_viz_img: bool = False,
) -> torch.Tensor:
url = f"https://raw.githubusercontent.com/naver/croco/d3d0ab2858d44bcad54e5bfc24f565983fbe18d9/assets/Chateau{img_selection}.png"
image = Image.open(requests.get(url, stream=True).raw)
image = image.resize(image_size)
image = image.convert("RGB")
img = torch.from_numpy(np.array(image))
viz_img = img.clone()
# Normalize the image
image_normalization = IMAGE_NORMALIZATION_DICT[image_norm_type]
img_mean = image_normalization.mean
img_std = image_normalization.std
img = (img.float() / 255.0 - img_mean) / img_std
# Convert to BCHW format
img = img.permute(2, 0, 1).unsqueeze(0).to(device)
if return_viz_img:
return img, viz_img
else:
return img
return _get_example_input
def inference_encoder(encoder, encoder_input):
# Encoder expects a ViTEncoderInput object
return encoder(encoder_input).features
def test_make_dummy_encoder(device):
print(f"Testing Init of Dummy Encoder on {device.type.upper()}")
encoder = _make_encoder_test("dummy").to(device)
# Check if the encoder has parameters
try:
params = list(encoder.parameters())
if not params:
print("Warning: The encoder has no parameters.")
else:
# Verify if the model is on the right device
assert params[0].is_cuda == (device.type == "cuda")
except Exception as e:
print(f"Error: {e}")
assert False # Fail the test if any error occurs
assert encoder is not None
def test_all_encoder_basics(encoders, encoder_configs, norm_types, example_input, encoder_name, device):
if encoder_name:
encoders = [encoder_name] # Override default encoders with the one specified
for encoder_name, encoder_config in zip(encoders, encoder_configs):
print(f"Testing encoder: {encoder_name} on {device.type.upper()}")
encoder = _make_encoder_test(encoder_name, **encoder_config).to(device)
_check_baseclass_attribute(encoder, norm_types)
_check_norm_check_function(encoder)
if isinstance(encoder, UniCeptionViTEncoderBase):
_check_vit_encoder_attribute(encoder)
_test_vit_encoder_patch_size(encoder, example_input)
def _check_baseclass_attribute(encoder, norm_types):
assert hasattr(encoder, "name")
assert hasattr(encoder, "size")
assert hasattr(encoder, "data_norm_type")
assert isinstance(encoder.name, str)
assert isinstance(encoder.size, str) or encoder.size is None
assert isinstance(encoder.data_norm_type, str)
# Check if the data_norm_type is in the list of normalization types
assert encoder.data_norm_type in norm_types
def _check_norm_check_function(encoder):
assert hasattr(encoder, "_check_data_normalization_type")
encoder_notm_type = encoder.data_norm_type
try:
encoder._check_data_normalization_type(encoder_notm_type)
except AssertionError:
assert False
try:
encoder._check_data_normalization_type("some_nonexistent_norm_type")
assert False
except AssertionError:
pass
def _check_vit_encoder_attribute(encoder):
assert hasattr(encoder, "patch_size")
assert isinstance(encoder.patch_size, int)
assert encoder.patch_size > 0
def _test_vit_encoder_patch_size(encoder, example_input):
print(f"Testing {encoder.name} inference")
image_size = (14 * encoder.patch_size, 14 * encoder.patch_size)
img = example_input(image_size, encoder.data_norm_type)
# Create an instance of ViTEncoderInput with correct attributes
encoder_input = ViTEncoderInput(
data_norm_type=encoder.data_norm_type,
image=img,
)
encoder_output = inference_encoder(encoder, encoder_input)
assert isinstance(encoder_output, torch.Tensor)
assert encoder_output.shape[2] == 14
assert encoder_output.shape[3] == 14
@pytest.fixture(scope="session", autouse=True)
def seed_everything():
seed = 42
random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
print(f"Seed set to: {seed} (type: {type(seed)})")
# Turn XFormers off for testing on CPU
os.environ["XFORMERS_DISABLED"] = "1"
|