File size: 2,694 Bytes
3ea5987 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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) |