init project

modules/pe3r/images.py

@@ -53,32 +53,32 @@ class Images:
             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)
+        # -- 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, :, :, :]
+        # -- 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:]))
+            self.sam1_images.append(transformed_image)
+            self.sam1_images_size.append(tuple(transformed_image.shape[-2:]))