Adding presets in the object detection reference scripts.

Author: datumbox

references/detection/presets.py

@@ -0,0 +1,21 @@
+import transforms as T
+
+
+class DetectionPresetTrain:
+    def __init__(self, hflip_prob=0.5):
+        trans = [T.ToTensor()]
+        if hflip_prob > 0:
+            trans.append(T.RandomHorizontalFlip(hflip_prob))
+
+        self.transforms = T.Compose(trans)
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)
+
+
+class DetectionPresetEval:
+    def __init__(self):
+        self.transforms = T.ToTensor()
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)

references/detection/train.py

@@ -32,8 +32,8 @@
 from group_by_aspect_ratio import GroupedBatchSampler, create_aspect_ratio_groups
 from engine import train_one_epoch, evaluate
 
+import presets
 import utils
-import transforms as T
 
 
 def get_dataset(name, image_set, transform, data_path):
@@ -48,11 +48,7 @@ def get_dataset(name, image_set, transform, data_path):
 
 
 def get_transform(train):
-    transforms = []
-    transforms.append(T.ToTensor())
-    if train:
-        transforms.append(T.RandomHorizontalFlip(0.5))
-    return T.Compose(transforms)
+    return presets.DetectionPresetTrain() if train else presets.DetectionPresetEval()
 
 
 def main(args):