Revert "work on arc data"

This reverts commit fcafbad.

configs/data.yaml

@@ -49,21 +49,21 @@ mnist:
     train: False
     download: True
 
-mnist-subset:
-  dir: mnist/
-  beton_dir: beton/mnist/
-  mean: [ 0.1307 ]
-  std: [ 0.3081 ]
-  num_classes: 10
-  num_channels: 1
-  img_width: 28
-  img_height: 28
-  train_dataset_params:
-    train: True
-    download: True
-  test_dataset_params:
-    train: False
-    download: True
+  mnist-subset:
+    dir: mnist/
+    beton_dir: beton/mnist/
+    mean: [ 0.1307 ]
+    std: [ 0.3081 ]
+    num_classes: 10
+    num_channels: 1
+    img_width: 28
+    img_height: 28
+    train_dataset_params:
+      train: True
+      download: True
+    test_dataset_params:
+      train: False
+      download: True
 
 imagenet:
   dir: imagenet/
@@ -152,13 +152,6 @@ eight_bit_numbers:
     samples_per_class: 1
     include_noise: False
 
-arc:
-  num_channels: 1
-  mean: [ 0. ]
-  std: [ 1. ]
-  img_width: 32
-  img_height: 32
-
 augmentation_v1:
   aug1:
     - probability: 0.8

src/data/from_conf.py

@@ -9,7 +9,6 @@
 from data.augmentation import get_image_augmentation, get_transform_pipeline
 from data.custom_datasets.eight_bit_numbers import EightBitDataset
 from data.custom_datasets.straight_line import StraightLine
-from data.custom_datasets.arc import ArcDataset
 from data.loader import get_ffcv_data_loaders, get_ffcv_image_pipeline, get_ffcv_label_pipeline, \
     get_torch_data_loaders
 
@@ -60,10 +59,6 @@ def _get_dataset(
         train_set = EightBitDataset(transform=transform, **dataset_config['train_dataset_params'])
         valid_set = EightBitDataset(transform=transform, **dataset_config['valid_dataset_params'])
         test_set = EightBitDataset(transform=transform, **dataset_config['test_dataset_params'])
-    elif dataset_name == "arc":
-        train_set = ArcDataset()
-        valid_set = None
-        test_set = ArcDataset()
     else:
         raise ValueError("Unknown dataset name: {}".format(dataset_name))
 

src/data/utils/plot_image.py

@@ -6,7 +6,6 @@
 import torchvision.transforms as T
 from PIL.Image import Image
 from mpl_toolkits.axes_grid1 import make_axes_locatable
-import matplotlib.colors as colors
 
 def undo_norm(img: torch.Tensor, mean: torch.Tensor, std: torch.Tensor):
     """
@@ -131,11 +130,7 @@ def plot_images(
             fig.colorbar(im, cax=cax, orientation='vertical')
 
         if masks is not None and mask is not None:
-            cmap = plt.cm.get_cmap('tab20')
-            cmaplist = [cmap(i) for i in range(cmap.N)]
-            cmaplist[0] = (0, 0, 0, 1.0)
-            cmap = colors.LinearSegmentedColormap.from_list(f'tab20_modified', cmaplist, cmap.N)
-            ax.imshow(mask, alpha=0.7, cmap=cmap, interpolation=interpolation_, vmin=mask_vmin, vmax=mask_vmax)
+            ax.imshow(mask, alpha=0.6, cmap='jet', interpolation=interpolation_, vmin=mask_vmin, vmax=mask_vmax)
 
         if lbl is not None:
             lbl = str(lbl.item()) if isinstance(lbl, torch.Tensor) else lbl