|
import numpy as np |
|
import torch |
|
from PIL import Image |
|
from modules.mobilesamv2.utils.transforms import ResizeLongestSide |
|
|
|
from modules.dust3r.utils.image import _resize_pil_image |
|
|
|
class Images: |
|
def __init__(self, filelist, device, size=512): |
|
|
|
self.pil_images = [] |
|
self.pil_images_size = [] |
|
self.np_images = [] |
|
self.np_images_size = [] |
|
|
|
tmp_images = [] |
|
first_image_size = None |
|
all_images_same_size = True |
|
for img_path in filelist: |
|
pil_image = Image.open(img_path).convert("RGB") |
|
tmp_images.append(pil_image) |
|
|
|
current_image_size = pil_image.size |
|
if first_image_size is None: |
|
first_image_size = current_image_size |
|
else: |
|
if current_image_size != first_image_size: |
|
all_images_same_size = False |
|
|
|
for img in tmp_images: |
|
|
|
if not all_images_same_size: |
|
|
|
pil_image = _resize_pil_image(img, size) |
|
W, H = pil_image.size |
|
cx, cy = W//2, H//2 |
|
halfw, halfh = ((2*cx)//16)*8, ((2*cy)//16)*8 |
|
if W == H: |
|
halfh = 3*halfw/4 |
|
pil_image = pil_image.crop((cx-halfw, cy-halfh, cx+halfw, cy+halfh)) |
|
else: |
|
pil_image = img |
|
|
|
np_image = np.array(pil_image) |
|
|
|
height, width = pil_image.size |
|
np_shape = np_image.shape[:2] |
|
|
|
self.pil_images.append(pil_image) |
|
self.np_images.append(np_image) |
|
|
|
self.pil_images_size.append((height, width)) |
|
self.np_images_size.append(np_shape) |
|
|
|
|
|
|
|
img_mean = torch.tensor((0.485, 0.456, 0.406))[:, None, None] |
|
img_std = torch.tensor((0.229, 0.224, 0.225))[:, None, None] |
|
self.sam2_images = [] |
|
|
|
self.sam2_video_size = (self.pil_images_size[0][1], self.pil_images_size[0][0]) |
|
self.sam2_input_size = 512 |
|
for pil_image in self.pil_images: |
|
np_image = np.array(pil_image.resize((self.sam2_input_size, self.sam2_input_size))) |
|
np_image = np_image / 255.0 |
|
sam2_image = torch.from_numpy(np_image).permute(2, 0, 1) |
|
self.sam2_images.append(sam2_image) |
|
self.sam2_images = torch.stack(self.sam2_images) |
|
self.sam2_images -= img_mean |
|
self.sam2_images /= img_std |
|
self.sam2_images.to(device) |
|
|
|
|
|
self.sam1_images = [] |
|
self.sam1_images_size = [] |
|
self.sam1_input_size = 1024 |
|
self.sam1_transform = ResizeLongestSide(self.sam1_input_size) |
|
for np_image in self.np_images: |
|
sam1_image = self.sam1_transform.apply_image(np_image) |
|
sam1_image_torch = torch.as_tensor(sam1_image, device=device) |
|
transformed_image = sam1_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :] |
|
|
|
self.sam1_images.append(transformed_image) |
|
self.sam1_images_size.append(tuple(transformed_image.shape[-2:])) |
