From d5d248594d2ba1239df10f3a7f1e2478ba52f62c Mon Sep 17 00:00:00 2001
From: Isaac Corley <22203655+isaaccorley@users.noreply.github.com>
Date: Thu, 2 Sep 2021 17:02:11 -0500
Subject: [PATCH] added whu-rs19 dataset, datamodule, and download script.
 updated readme

---
 scripts/download_whurs19.sh           |  4 ++++
 torchrs/datasets/whu_rs19.py          | 20 ++++++++++++++++++++
 torchrs/train/datamodules/whu_rs19.py | 26 ++++++++++++++++++++++++++
 3 files changed, 50 insertions(+)
 create mode 100644 scripts/download_whurs19.sh
 create mode 100644 torchrs/datasets/whu_rs19.py
 create mode 100644 torchrs/train/datamodules/whu_rs19.py

diff --git a/scripts/download_whurs19.sh b/scripts/download_whurs19.sh
new file mode 100644
index 0000000..f560023
--- /dev/null
+++ b/scripts/download_whurs19.sh
@@ -0,0 +1,4 @@
+mkdir -p .data/
+wget https://github.com/CAPTAIN-WHU/BED4RS/raw/main/datasets/WHU-RS19.zip
+unzip WHU-RS19.zip -d .data/
+rm WHU-RS19.zip
diff --git a/torchrs/datasets/whu_rs19.py b/torchrs/datasets/whu_rs19.py
new file mode 100644
index 0000000..99db425
--- /dev/null
+++ b/torchrs/datasets/whu_rs19.py
@@ -0,0 +1,20 @@
+import os
+
+import torchvision.transforms as T
+from torchvision.datasets import ImageFolder
+
+
+class WHU_RS19(ImageFolder):
+    """ WHU-RS19 dataset from'Structural High-resolution Satellite Image Indexing', Xia at al. (2010)
+    https://hal.archives-ouvertes.fr/file/index/docid/458685/filename/structural_satellite_indexing_XYDG.pdf
+
+    """
+    def __init__(
+        self,
+        root: str = ".data/WHU-RS19",
+        transform: T.Compose = T.Compose([T.ToTensor()])
+    ):
+        super().__init__(
+            root=root,
+            transform=transform
+        )
diff --git a/torchrs/train/datamodules/whu_rs19.py b/torchrs/train/datamodules/whu_rs19.py
new file mode 100644
index 0000000..427ca35
--- /dev/null
+++ b/torchrs/train/datamodules/whu_rs19.py
@@ -0,0 +1,26 @@
+from typing import Optional
+
+import torchvision.transforms as T
+
+from torchrs.datasets.utils import dataset_split
+from torchrs.train.datamodules import BaseDataModule
+from torchrs.datasets import WHU_RS19
+
+
+class WHU_RS19DataModule(BaseDataModule):
+
+    def __init__(
+        self,
+        root: str = ".data/WHU-RS19",
+        transform: T.Compose = T.Compose([T.ToTensor()]),
+        *args, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.root = root
+        self.transform = transform
+
+    def setup(self, stage: Optional[str] = None):
+        dataset = WHU_RS19(root=self.root, transform=self.transform)
+        self.train_dataset, self.val_dataset, self.test_dataset = dataset_split(
+            dataset, val_pct=self.val_split, test_pct=self.test_split
+        )