File size: 2,431 Bytes
c881b77 | 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 | from . import clip
import torch
import numpy as np
import cv2
_CONTOUR_INDEX = 1 if cv2.__version__.split('.')[0] == '3' else 0
class ClipOutputTarget:
def __init__(self, category):
self.category = category
def __call__(self, model_output):
if len(model_output.shape) == 1:
return model_output[self.category]
return model_output[:, self.category]
def reshape_transform(tensor, height=28, width=28):
tensor = tensor.permute(1, 0, 2)
result = tensor[:, 1:, :].reshape(tensor.size(0), height, width, tensor.size(2))
# Bring the channels to the first dimension,
# like in CNNs.
result = result.transpose(2, 3).transpose(1, 2)
return result
def zeroshot_classifier(classnames, templates, model, device):
with torch.no_grad():
zeroshot_weights = []
for classname in classnames:
texts = [template.format(classname) for template in templates] #format with class
texts = clip.tokenize(texts).to(device) #tokenize
class_embeddings = model.encode_text(texts) #embed with text encoder
class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
class_embedding = class_embeddings.mean(dim=0)
class_embedding /= class_embedding.norm()
zeroshot_weights.append(class_embedding)
zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
return zeroshot_weights.t()
def scoremap2bbox(scoremap, threshold, multi_contour_eval=False):
height, width = scoremap.shape
scoremap_image = np.expand_dims((scoremap * 255).astype(np.uint8), 2)
_, thr_gray_heatmap = cv2.threshold(
src=scoremap_image,
thresh=int(threshold * np.max(scoremap_image)),
maxval=255,
type=cv2.THRESH_BINARY)
contours = cv2.findContours(
image=thr_gray_heatmap,
mode=cv2.RETR_TREE,
method=cv2.CHAIN_APPROX_SIMPLE)[_CONTOUR_INDEX]
if len(contours) == 0:
return np.asarray([[0, 0, 0, 0]]), 1
if not multi_contour_eval:
contours = [max(contours, key=cv2.contourArea)]
estimated_boxes = []
for contour in contours:
x, y, w, h = cv2.boundingRect(contour)
x0, y0, x1, y1 = x, y, x + w, y + h
x1 = min(x1, width - 1)
y1 = min(y1, height - 1)
estimated_boxes.append([x0, y0, x1, y1])
return np.asarray(estimated_boxes), len(contours) |