Skip to content

Commit

Permalink
added whu-rs19 dataset, datamodule, and download script. updated readme
Browse files Browse the repository at this point in the history
  • Loading branch information
isaaccorley committed Sep 2, 2021
1 parent f317341 commit d5d2485
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 0 deletions.
4 changes: 4 additions & 0 deletions scripts/download_whurs19.sh
Original file line number Diff line number Diff line change
@@ -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
20 changes: 20 additions & 0 deletions torchrs/datasets/whu_rs19.py
Original file line number Diff line number Diff line change
@@ -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
)
26 changes: 26 additions & 0 deletions torchrs/train/datamodules/whu_rs19.py
Original file line number Diff line number Diff line change
@@ -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
)

0 comments on commit d5d2485

Please sign in to comment.