Views
No views yet
1import torch
2import cv2
3import numpy as np
4from torchvision import transforms
5from transformers import AutoImageProcessor, AutoModel
6
7def padding(img):
8 """
9 :param img: take the image as input, np.uint8 format, [0-255] range
10 :return: img: square image with white pixel padded along the shorter side.
11 """
12 h, w, _ = img.shape
13 if h > w:
14 new_img = 255 * np.ones((h, h, 3)).astype(np.uint8)
15 start_w = int((h-w)/2)
16 new_img[:, start_w:start_w+w, :] = img
17 return new_img
18
19 elif h < w:
20 new_img = 255 * np.ones((w, w, 3)).astype(np.uint8)
21 start_h = int((w - h) / 2)
22 new_img[start_h:start_h + h, :, :] = img
23 return new_img
24 else:
25 return img
26
27image_path = "your local image path"
28dim = 384
29image = cv2.imread(image_path)
30image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
31image = padding(image)
32image = cv2.resize(image, (dim, dim))
33transform = transforms.ToTensor()
34image = transform(image)
35image = image.unsqueeze(0)
36model_name = "zhanganyi88/Cat2Real-DINOv3-384"
37processor = AutoImageProcessor.from_pretrained(model_name)
38model = AutoModel.from_pretrained(
39 model_name,
40 device_map="auto",
41)
42inputs = processor(images=image, return_tensors="pt").to(model.device)
43with torch.inference_mode():
44 outputs = model(**inputs)
45embedding = outputs.pooler_output
46print("embedding shape:", embedding.shape)