Refactor code with black

BDD_OIA/BDD/data_processing.py

@@ -3,14 +3,14 @@
 Each dataset is a list of metadata, each includes official image id, full image path, class label, attribute labels, attribute certainty scores, and attribute labels calibrated for uncertainty
 """
 
-import sys
+import argparse
 import os
-import random
 import pickle
-import argparse
-from os import listdir
-from os.path import isfile, isdir, join
+import random
+import sys
 from collections import defaultdict as ddict
+from os import listdir
+from os.path import isdir, isfile, join
 
 
 def extract_data(data_dir):
@@ -44,20 +44,38 @@ def extract_data(data_dir):
         },  # calibrate main label based on uncertainty label
         0: {1: 0, 2: 0.5, 3: 0.25, 4: 0},
     }
-    with open(join(cwd, data_dir + "/attributes/image_attribute_labels.txt"), "r") as f:
+    with open(
+        join(
+            cwd, data_dir + "/attributes/image_attribute_labels.txt"
+        ),
+        "r",
+    ) as f:
         for line in f:
-            file_idx, attribute_idx, attribute_label, attribute_certainty = (
-                line.strip().split()[:4]
-            )
+            (
+                file_idx,
+                attribute_idx,
+                attribute_label,
+                attribute_certainty,
+            ) = line.strip().split()[:4]
             attribute_label = int(attribute_label)
             attribute_certainty = int(attribute_certainty)
-            uncertain_label = uncertainty_map[attribute_label][attribute_certainty]
-            attribute_labels_all[int(file_idx)].append(attribute_label)
-            attribute_uncertain_labels_all[int(file_idx)].append(uncertain_label)
-            attribute_certainties_all[int(file_idx)].append(attribute_certainty)
+            uncertain_label = uncertainty_map[attribute_label][
+                attribute_certainty
+            ]
+            attribute_labels_all[int(file_idx)].append(
+                attribute_label
+            )
+            attribute_uncertain_labels_all[int(file_idx)].append(
+                uncertain_label
+            )
+            attribute_certainties_all[int(file_idx)].append(
+                attribute_certainty
+            )
 
     is_train_test = dict()  # map from image id to 0 / 1 (1 = train)
-    with open(join(cwd, data_dir + "/train_test_split.txt"), "r") as f:
+    with open(
+        join(cwd, data_dir + "/train_test_split.txt"), "r"
+    ) as f:
         for line in f:
             idx, is_train = line.strip().split()
             is_train_test[int(idx)] = int(is_train)
@@ -68,7 +86,9 @@ def extract_data(data_dir):
 
     train_val_data, test_data = [], []
     train_data, val_data = [], []
-    folder_list = [f for f in listdir(data_path) if isdir(join(data_path, f))]
+    folder_list = [
+        f for f in listdir(data_path) if isdir(join(data_path, f))
+    ]
     folder_list.sort()  # sort by class index
     for i, folder in enumerate(folder_list):
         folder_path = join(data_path, folder)
@@ -86,8 +106,12 @@ def extract_data(data_dir):
                 "img_path": img_path,
                 "class_label": i,
                 "attribute_label": attribute_labels_all[img_id],
-                "attribute_certainty": attribute_certainties_all[img_id],
-                "uncertain_attribute_label": attribute_uncertain_labels_all[img_id],
+                "attribute_certainty": attribute_certainties_all[
+                    img_id
+                ],
+                "uncertain_attribute_label": attribute_uncertain_labels_all[
+                    img_id
+                ],
             }
             if is_train_test[img_id]:
                 train_val_data.append(metadata)
@@ -107,9 +131,15 @@ def extract_data(data_dir):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Dataset preparation")
-    parser.add_argument("--save_dir", "-d", help="Where to save the new datasets")
-    parser.add_argument("--data_dir", help="Where to load the datasets")
+    parser = argparse.ArgumentParser(
+        description="Dataset preparation"
+    )
+    parser.add_argument(
+        "--save_dir", "-d", help="Where to save the new datasets"
+    )
+    parser.add_argument(
+        "--data_dir", help="Where to load the datasets"
+    )
     args = parser.parse_args()
     print(args.data_dir)
     train_data, val_data, test_data = extract_data(args.data_dir)

BDD_OIA/BDD/dataset.py

@@ -2,17 +2,16 @@
 General utils for training, evaluation and data loading
 """
 
-import sys
 import os
-import torch
 import pickle
+import sys
+
 import numpy as np
+import torch
 import torchvision.transforms as transforms
-
-from PIL import Image
 from BDD.config import BASE_DIR, N_ATTRIBUTES
-from torch.utils.data import BatchSampler
-from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+from torch.utils.data import BatchSampler, DataLoader, Dataset
 
 
 class BDDDataset(Dataset):
@@ -41,9 +40,16 @@ def __init__(
         transform: whether to apply any special transformation. Default = None, i.e. use standard ImageNet preprocessing
         """
         self.data = []
-        self.is_train = any(["train" in path for path in pkl_file_paths])
+        self.is_train = any(
+            ["train" in path for path in pkl_file_paths]
+        )
         if not self.is_train:
-            assert any([("test" in path) or ("val" in path) for path in pkl_file_paths])
+            assert any(
+                [
+                    ("test" in path) or ("val" in path)
+                    for path in pkl_file_paths
+                ]
+            )
         for file_path in pkl_file_paths:
             self.data.extend(pickle.load(open(file_path, "rb")))
         self.transform = transform
@@ -133,7 +139,9 @@ class ImbalancedDatasetSampler(torch.utils.data.sampler.Sampler):
     def __init__(self, dataset, indices=None):
         # if indices is not provided,
         # all elements in the dataset will be considered
-        self.indices = list(range(len(dataset))) if indices is None else indices
+        self.indices = (
+            list(range(len(dataset))) if indices is None else indices
+        )
 
         # if num_samples is not provided,
         # draw `len(indices)` samples in each iteration
@@ -150,17 +158,22 @@ def __init__(self, dataset, indices=None):
 
         # weight for each sample
         weights = [
-            1.0 / label_to_count[self._get_label(dataset, idx)] for idx in self.indices
+            1.0 / label_to_count[self._get_label(dataset, idx)]
+            for idx in self.indices
         ]
         self.weights = torch.DoubleTensor(weights)
 
-    def _get_label(self, dataset, idx):  # Note: for single attribute dataset
+    def _get_label(
+        self, dataset, idx
+    ):  # Note: for single attribute dataset
         return dataset.data[idx]["attribute_label"][0]
 
     def __iter__(self):
         idx = (
             self.indices[i]
-            for i in torch.multinomial(self.weights, self.num_samples, replacement=True)
+            for i in torch.multinomial(
+                self.weights, self.num_samples, replacement=True
+            )
         )
         return idx
 
@@ -218,7 +231,10 @@ def load_data(
         loader = DataLoader(dataset, batch_sampler=sampler)
     else:
         loader = DataLoader(
-            dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last
+            dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            drop_last=drop_last,
         )
     return loader
 

BDD_OIA/BDD/template_model.py

@@ -1,11 +1,12 @@
+import math
 import os
+
 import torch
 import torch.nn as nn
-from torch.nn import Parameter
 import torch.nn.functional as F
 import torch.utils.model_zoo as model_zoo
 import torchvision
-import math
+from torch.nn import Parameter
 
 __all__ = ["FRCNN", "RCNN_global"]
 
@@ -28,7 +29,9 @@ def __init__(self, cfg=None, random_select=False):
         self.conv_glob2 = nn.Conv2d(128, 64, 3, padding=1)
         self.relu_glob2 = nn.ReLU(inplace=True)
 
-        self.lin_glob = nn.Linear(in_features=3136, out_features=2048, bias=True)
+        self.lin_glob = nn.Linear(
+            in_features=3136, out_features=2048, bias=True
+        )
         self.relu_glob = nn.ReLU()
 
     def forward(self, x):