Update modeling_visual.py

modeling_visual.py

@@ -22,68 +22,6 @@ from transformers import AutoModel, AutoProcessor
 from transformers.activations import ACT2FN
 
 
-class TransformCXR(object):
-    def __init__(
-            self,
-            image_size=448,
-            mean=(0.48145466, 0.4578275, 0.40821073),
-            std=(0.26862954, 0.26130258, 0.27577711),
-            allow_shift=True,
-            training=True,
-            normalize=True
-    ):
-
-        resize_size = image_size
-        p_train = 0.5
-        shift_limit = (-0.0, 0.0)
-        scale_limit = (-0.1, -0.02)
-        rotate_limit = 5
-        scale = (0.00, 0.01)
-        brightness_limit = (-0.15, 0.15)
-        contrast_limit = (-0.05, 0.05)
-        pad_mode = cv2.BORDER_CONSTANT
-        pad_val = (0, 0, 0)
-
-        if training:
-            if allow_shift:
-                transform_list = [
-                    A.ShiftScaleRotate(
-                        shift_limit=shift_limit, scale_limit=scale_limit,
-                        rotate_limit=rotate_limit, border_mode=pad_mode, value=pad_val,
-                        p=p_train
-                    ),
-                    A.Perspective(
-                        scale=scale, pad_mode=pad_mode, pad_val=pad_val, p=p_train
-                    ),
-                    A.Resize(height=resize_size, width=resize_size, interpolation=cv2.INTER_CUBIC),
-                    A.RandomCrop(height=image_size, width=image_size),
-                    A.RandomBrightnessContrast(
-                        brightness_limit=brightness_limit, contrast_limit=contrast_limit,
-                        p=p_train
-                    )
-                ]
-            else:
-                transform_list = [
-                    A.Resize(height=image_size, width=image_size, interpolation=cv2.INTER_CUBIC),
-                    A.RandomBrightnessContrast(
-                        brightness_limit=brightness_limit, contrast_limit=contrast_limit,
-                        p=p_train
-                    )
-                ]
-        else:
-            transform_list = [
-                A.Resize(height=image_size, width=image_size, interpolation=cv2.INTER_CUBIC)
-            ]
-
-        if normalize:
-            transform_list += [A.Normalize(mean=mean, std=std), ToTensorV2(transpose_mask=True)]
-
-        self.transforms = A.Compose(transform_list)
-
-    def __call__(self, image):
-        image = np.array(image)
-        return self.transforms(image=image)['image']
-
 
 def get_abs_pos(abs_pos, tgt_size):
     # abs_pos: L, C
@@ -247,11 +185,7 @@ class CLIPModel(nn.Module):
         # Transforms
         self.mean = self.processor.image_mean
         self.std = self.processor.image_std
-        self.image_transform_train = TransformCXR(image_size=image_size, mean=self.mean, std=self.std, training=True)
-        self.image_transform_train_no_shift = TransformCXR(
-            image_size=image_size, mean=self.mean, std=self.std, allow_shift=False, training=True
-        )
-        self.image_transform_val = TransformCXR(image_size=image_size, mean=self.mean, std=self.std, training=False)
+        
         self.image_transform = transforms.Compose([
             transforms.Resize(
                 (image_size, image_size),
@@ -298,17 +232,7 @@ class CLIPModel(nn.Module):
 
         image = image.convert("RGB")
 
-        no_shift = any([keyword in image_path for keyword in ["vindr", "candid", "siim", "object-cxr", "ms-cxr"]])
-        try:
-            if training or self.training:
-                if no_shift:
-                    image_tensor = self.image_transform_train_no_shift(image)
-                else:
-                    image_tensor = self.image_transform_train(image)
-            else:
-                image_tensor = self.image_transform_val(image)
-        except:
-            image_tensor = self.image_transform(image)
+        image_tensor = self.image_transform(image)
         return image_tensor
 
     def encode(self, image_paths: List[str], training):