| import os |
|
|
| import torch |
| import wilds |
| from wilds.common.data_loaders import get_eval_loader, get_train_loader |
|
|
|
|
| class FMOW: |
| test_subset = None |
|
|
| def __init__(self, |
| preprocess, |
| location=os.path.expanduser('~/data'), |
| batch_size=128, |
| num_workers=16, |
| subset='test', |
| classnames=None, |
| **kwargs): |
|
|
| self.dataset = wilds.get_dataset(dataset='fmow', root_dir=location) |
|
|
| self.train_dataset = self.dataset.get_subset('train', transform=preprocess) |
| self.train_loader = get_train_loader("standard", self.train_dataset, num_workers=num_workers, batch_size=batch_size) |
|
|
| self.test_dataset = self.dataset.get_subset(self.test_subset, transform=preprocess) |
| self.test_loader = get_eval_loader("standard", self.test_dataset, num_workers=num_workers, batch_size=batch_size) |
|
|
| self.classnames = [ |
| "airport", "airport_hangar", "airport_terminal", "amusement_park", "aquaculture", |
| "archaeological_site", "barn", "border_checkpoint", "burial_site", "car_dealership", |
| "construction_site", "crop_field", "dam", "debris_or_rubble", "educational_institution", |
| "electric_substation", "factory_or_powerplant", "fire_station", "flooded_road", "fountain", |
| "gas_station", "golf_course", "ground_transportation_station", "helipad", "hospital", |
| "impoverished_settlement", "interchange", "lake_or_pond", "lighthouse", "military_facility", |
| "multi-unit_residential", "nuclear_powerplant", "office_building", "oil_or_gas_facility", "park", |
| "parking_lot_or_garage", "place_of_worship", "police_station", "port", "prison", "race_track", |
| "railway_bridge", "recreational_facility", "road_bridge", "runway", "shipyard", "shopping_mall", |
| "single-unit_residential", "smokestack", "solar_farm", "space_facility", "stadium", "storage_tank", |
| "surface_mine", "swimming_pool", "toll_booth", "tower", "tunnel_opening", "waste_disposal", |
| "water_treatment_facility", "wind_farm", "zoo" |
| ] |
|
|
| def post_loop_metrics(self, labels, preds, metadata, args): |
| metadata = torch.stack(metadata) |
| preds = preds.argmax(dim=1, keepdim=True).view_as(labels) |
| results = self.dataset.eval(preds, labels, metadata) |
| return results[0] |
|
|
| class FMOWID(FMOW): |
| def __init__(self, *args, **kwargs): |
| self.test_subset = 'id_test' |
| super().__init__(*args, **kwargs) |
|
|
| class FMOWOOD(FMOW): |
| def __init__(self, *args, **kwargs): |
| self.test_subset = 'test' |
| super().__init__(*args, **kwargs) |
