first commit

ssd/data/transforms/__init__.py

@@ -0,0 +1,34 @@
+from ssd.modeling.anchors.prior_box import PriorBox
+from .target_transform import SSDTargetTransform
+from .transforms import *
+
+
+def build_transforms(cfg, is_train=True):
+    if is_train:
+        transform = [
+            ConvertFromInts(),
+            PhotometricDistort(),
+            Expand(cfg.INPUT.PIXEL_MEAN),
+            RandomSampleCrop(),
+            RandomMirror(),
+            ToPercentCoords(),
+            Resize(cfg.INPUT.IMAGE_SIZE),
+            SubtractMeans(cfg.INPUT.PIXEL_MEAN),
+            ToTensor(),
+        ]
+    else:
+        transform = [
+            Resize(cfg.INPUT.IMAGE_SIZE),
+            SubtractMeans(cfg.INPUT.PIXEL_MEAN),
+            ToTensor()
+        ]
+    transform = Compose(transform)
+    return transform
+
+
+def build_target_transform(cfg):
+    transform = SSDTargetTransform(PriorBox(cfg)(),
+                                   cfg.MODEL.CENTER_VARIANCE,
+                                   cfg.MODEL.SIZE_VARIANCE,
+                                   cfg.MODEL.THRESHOLD)
+    return transform