Spaces:
Running on Zero
Running on Zero
File size: 5,322 Bytes
c7b663e | 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 | import numpy as np
import torch
import torch.nn.functional as F
from einops import rearrange
from tqdm import tqdm
# from open_clip import get_cast_dtype
# from training.precision import get_autocast
from .voc_dataset import voc_extended_classes
import os
import re
import scipy
import numpy as np
import time
from eval.models import dust3r, mast3r, Sab3r
from datasets import load_dataset
import torch
import torchvision.transforms as T
from PIL import Image
import PIL.Image
import PIL
from featup.util import norm, unnorm
from featup.plotting import plot_feats, plot_lang_heatmaps
import torchvision.transforms as tvf
from pytorch_lightning import seed_everything
from featup.util import pca, remove_axes
import matplotlib.pyplot as plt
from featup.featurizers.maskclip.clip import tokenize
from tqdm import tqdm
class PascalMIoU:
def __init__(self):
self._num_classes = 20 + 1 # background
self.confusion_matrix = None
self._num_examples = 0
self.reset()
def reset(self):
self.confusion_matrix = np.zeros((self._num_classes,) * 2, dtype=np.int64)
self._num_examples = 0
def update(self, prediction, label):
self._num_examples += label.shape[0]
prediction[prediction >= self._num_classes] = 0
mask = label < 255 # skip-pixel for VOC dataset
indices = self._num_classes * label[mask] + prediction[mask]
m = np.bincount(indices, minlength=self._num_classes**2).reshape(
self._num_classes, self._num_classes
)
self.confusion_matrix += m
def compute(self):
if self._num_examples == 0:
return 0
cm = self.confusion_matrix
iou_list = cm.diagonal() / (
cm.sum(axis=0) + cm.sum(axis=1) - cm.diagonal() + np.finfo(np.float32).eps
)
return {"mIoU": np.nanmean(iou_list), "per_class": iou_list.tolist()}
def get_similarity(
image_encodings,
label_encodings,
target_shape,
interpolation="bilinear",
do_argmax=False,
):
"""
Args:
image_encodings:
label_encodings:
target_shape:
interpolation: nearest, bilinear
do_argmax:
Returns:
"""
image_encodings = image_encodings.cpu()
label_encodings = label_encodings.cpu()
image_encodings = rearrange(image_encodings, "b d h w -> b h w d")
similarity = image_encodings @ label_encodings.T
similarity = rearrange(similarity, "b h w d-> b d h w")
if do_argmax:
similarity = torch.argmax(similarity, dim=1, keepdim=True).to(torch.float64)
return similarity
class PascalEvaluator:
def __init__(self, model, upsampler, dataset, metric, opts):
self.classes = voc_extended_classes
self._class_prompts = self.get_class_prompts()
self.class_embeddings = None
self.model = model
self.upsampler = upsampler
self.dataset = dataset
self.metric: PascalMIoU = metric
self.opts = opts
self.dataloader = torch.utils.data.DataLoader(
dataset, batch_size=opts.batch_size, num_workers=opts.workers
)
def evaluate(self):
self.metric.reset()
if self.class_embeddings is None:
self.class_embeddings = self.get_class_embeddings()
# autocast = get_autocast(self.opts.precision)
# cast_dtype = get_cast_dtype(self.opts.precision)
with torch.no_grad():
for images, target in tqdm(
self.dataloader, unit_scale=self.opts.batch_size
):
breakpoint()
image_features = self.model.predict(images)
clip_features = (
image_features.get_clip()
) # Assuming this returns features of shape (512, H, W)
pred_0 = (
torch.tensor(clip_features[0]).cuda().contiguous().permute(2, 0, 1)
).unsqueeze(
0
) # Shape should be (1, 512, H, W)
similarity = get_similarity(
pred_0, self.class_embeddings, target.shape, do_argmax=True
)
similarity = similarity[:, 0, :, :]
pred = similarity.detach().to(torch.int64) # .to(self.opts.device)
target = target.to(torch.int64)
# target = target.to(self.opts.device)
self.metric.update(pred, target)
def get_class_embeddings(self):
# self.model.eval()
cls_names = [name.lower() for name in self._class_prompts.values()]
with torch.no_grad():
tokenized_text = tokenize(cls_names).to(
"cuda"
) # Tokenize and move to the correct device
class_embeddings = self.upsampler.model.model.encode_text(tokenized_text)
class_embeddings = F.normalize(class_embeddings, dim=-1)
class_embeddings /= class_embeddings.norm()
return class_embeddings
def get_class_prompts(self):
class_prompts = {}
for idx, c in enumerate(self.classes):
if c.startswith(tuple("aeiou")):
class_prompts[idx] = f"a photo of an {c}"
else:
class_prompts[idx] = f"a photo of a {c}"
return class_prompts
|