| 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 = [] | |
| # -- original images -- | |
| 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: | |
| # resize long side to 512 | |
| 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) | |
| # # -- sam2 images -- | |
| # 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 = [] | |
| # # TODO | |
| # 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) | |
| # # -- sam1 images -- | |
| # 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:])) |