Upload 7 files

models/networks/heads/auxilliary.py

@@ -0,0 +1,33 @@
+import torch.nn as nn
+from models.networks.utils import UnormGPS
+from torch.nn.functional import tanh, sigmoid, softmax
+
+
+class AuxHead(nn.Module):
+    def __init__(self, aux_data=[], use_tanh=False):
+        super().__init__()
+        self.aux_data = aux_data
+        self.unorm = UnormGPS()
+        self.use_tanh = use_tanh
+
+    def forward(self, x):
+        """Forward pass of the network.
+        x : Union[torch.Tensor, dict] with the output of the backbone.
+        """
+        if self.use_tanh:
+            gps = tanh(x["gps"])
+        gps = self.unorm(gps)
+        output = {"gps": gps}
+        if "land_cover" in self.aux_data:
+            output["land_cover"] = softmax(x["land_cover"])
+        if "road_index" in self.aux_data:
+            output["road_index"] = x["road_index"]
+        if "drive_side" in self.aux_data:
+            output["drive_side"] = sigmoid(x["drive_side"])
+        if "climate" in self.aux_data:
+            output["climate"] = softmax(x["climate"])
+        if "soil" in self.aux_data:
+            output["soil"] = softmax(x["soil"])
+        if "dist_sea" in self.aux_data:
+            output["dist_sea"] = x["dist_sea"]
+        return output

models/networks/heads/classification.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+
+
+class ClassificationHead(nn.Module):
+    """Classification head for the network."""
+
+    def __init__(self, id_to_gps):
+        super().__init__()
+        self.id_to_gps = id_to_gps
+
+    def forward(self, x):
+        """Forward pass of the network.
+        x : Union[torch.Tensor, dict] with the output of the backbone.
+        """
+        gps = self.id_to_gps(x.argmax(dim=-1))
+        return {"label": x, **gps}

models/networks/heads/hybrid.py

@@ -0,0 +1,194 @@
+import torch
+import torch.nn as nn
+import pandas as pd
+
+from models.networks.utils import UnormGPS
+
+
+class HybridHead(nn.Module):
+    """Classification head followed by regression head for the network."""
+
+    def __init__(self, final_dim, quadtree_path, use_tanh, scale_tanh):
+        super().__init__()
+        self.final_dim = final_dim
+        self.use_tanh = use_tanh
+        self.scale_tanh = scale_tanh
+
+        self.unorm = UnormGPS()
+
+        if quadtree_path is not None:
+            quadtree = pd.read_csv(quadtree_path)
+            self.init_quadtree(quadtree)
+
+    def init_quadtree(self, quadtree):
+        quadtree[["min_lat", "max_lat"]] /= 90.0
+        quadtree[["min_lon", "max_lon"]] /= 180.0
+        self.register_buffer(
+            "cell_center",
+            0.5 * torch.tensor(quadtree[["max_lat", "max_lon"]].values)
+            + 0.5 * torch.tensor(quadtree[["min_lat", "min_lon"]].values),
+        )
+        self.register_buffer(
+            "cell_size",
+            torch.tensor(quadtree[["max_lat", "max_lon"]].values)
+            - torch.tensor(quadtree[["min_lat", "min_lon"]].values),
+        )
+
+    def forward(self, x, gt_label):
+        """Forward pass of the network.
+        x : Union[torch.Tensor, dict] with the output of the backbone.
+        """
+
+        classification_logits = x[..., : self.final_dim]
+        classification = classification_logits.argmax(dim=-1)
+
+        regression = x[..., self.final_dim :]
+
+        if self.use_tanh:
+            regression = self.scale_tanh * torch.tanh(regression)
+
+        regression = regression.view(regression.shape[0], -1, 2)
+
+        if self.training:
+            regression = torch.gather(
+                regression,
+                1,
+                gt_label.unsqueeze(-1).unsqueeze(-1).expand(regression.shape[0], 1, 2),
+            )[:, 0, :]
+            size = 2.0 / self.cell_size[gt_label]
+            center = self.cell_center[gt_label]
+            gps = (
+                self.cell_center[gt_label] + regression * self.cell_size[gt_label] / 2.0
+            )
+        else:
+            regression = torch.gather(
+                regression,
+                1,
+                classification.unsqueeze(-1)
+                .unsqueeze(-1)
+                .expand(regression.shape[0], 1, 2),
+            )[:, 0, :]
+            size = 2.0 / self.cell_size[classification]
+            center = self.cell_center[classification]
+            gps = (
+                self.cell_center[classification]
+                + regression * self.cell_size[classification] / 2.0
+            )
+
+        gps = self.unorm(gps)
+
+        return {
+            "label": classification_logits,
+            "gps": gps,
+            "size": size,
+            "center": center,
+            "reg": regression,
+        }
+
+class HybridHeadCentroid(nn.Module):
+    """Classification head followed by regression head for the network."""
+
+    def __init__(self, final_dim, quadtree_path, use_tanh, scale_tanh):
+        super().__init__()
+        self.final_dim = final_dim
+        self.use_tanh = use_tanh
+        self.scale_tanh = scale_tanh
+
+        self.unorm = UnormGPS()
+        if quadtree_path is not None:
+            quadtree = pd.read_csv(quadtree_path)
+            self.init_quadtree(quadtree)
+
+    def init_quadtree(self, quadtree):
+        quadtree[["min_lat", "max_lat", "mean_lat"]] /= 90.0
+        quadtree[["min_lon", "max_lon", "mean_lon"]] /= 180.0
+        self.cell_center = torch.tensor(quadtree[["mean_lat", "mean_lon"]].values)
+        self.cell_size_up = torch.tensor(quadtree[["max_lat", "max_lon"]].values) - torch.tensor(quadtree[["mean_lat", "mean_lon"]].values)
+        self.cell_size_down = torch.tensor(quadtree[["mean_lat", "mean_lon"]].values) - torch.tensor(quadtree[["min_lat", "min_lon"]].values)
+
+    def forward(self, x, gt_label):
+        """Forward pass of the network.
+        x : Union[torch.Tensor, dict] with the output of the backbone.
+        """
+        classification_logits = x[..., : self.final_dim]
+        classification = classification_logits.argmax(dim=-1)
+        self.cell_size_up = self.cell_size_up.to(classification.device)
+        self.cell_center = self.cell_center.to(classification.device)
+        self.cell_size_down = self.cell_size_down.to(classification.device)
+
+        regression = x[..., self.final_dim :]
+
+        if self.use_tanh:
+            regression = self.scale_tanh * torch.tanh(regression)
+
+        regression = regression.view(regression.shape[0], -1, 2)
+
+        if self.training:
+            regression = torch.gather(
+                regression,
+                1,
+                gt_label.unsqueeze(-1).unsqueeze(-1).expand(regression.shape[0], 1, 2),
+            )[:, 0, :]
+            size = torch.where(
+                regression > 0,
+                self.cell_size_up[gt_label],
+                self.cell_size_down[gt_label],
+            )
+            center = self.cell_center[gt_label]
+            gps = self.cell_center[gt_label] + regression * size
+        else:
+            regression = torch.gather(
+                regression,
+                1,
+                classification.unsqueeze(-1)
+                .unsqueeze(-1)
+                .expand(regression.shape[0], 1, 2),
+            )[:, 0, :]
+            size = torch.where(
+                regression > 0,
+                self.cell_size_up[classification],
+                self.cell_size_down[classification],
+            )
+            center = self.cell_center[classification]
+            gps = self.cell_center[classification] + regression * size
+
+        gps = self.unorm(gps)
+
+        return {
+            "label": classification_logits,
+            "gps": gps,
+            "size": 1.0 / size,
+            "center": center,
+            "reg": regression,
+        }
+
+
+class SharedHybridHead(HybridHead):
+    """Classification head followed by SHARED regression head for the network."""
+
+    def forward(self, x, gt_label):
+        """Forward pass of the network.
+        x : Union[torch.Tensor, dict] with the output of the backbone.
+        """
+
+        classification_logits = x[..., : self.final_dim]
+        classification = classification_logits.argmax(dim=-1)
+
+        regression = x[..., self.final_dim :]
+
+        if self.use_tanh:
+            regression = self.scale_tanh * torch.tanh(regression)
+
+        if self.training:
+            gps = (
+                self.cell_center[gt_label] + regression * self.cell_size[gt_label] / 2.0
+            )
+        else:
+            gps = (
+                self.cell_center[classification]
+                + regression * self.cell_size[classification] / 2.0
+            )
+
+        gps = self.unorm(gps)
+
+        return {"label": classification_logits, "gps": gps}

models/networks/heads/id_to_gps.py

@@ -0,0 +1,33 @@
+import torch
+from models.networks.utils import UnormGPS
+import torch.nn as nn
+import numpy as np
+
+
+class IdToGPS(nn.Module):
+    def __init__(self, id_to_gps: str):
+        """Map index to gps coordinates (indices can be country or city ids)"""
+        super().__init__()
+        if "quadtree" in id_to_gps:
+            self.id_to_gps = torch.load(
+                "_".join(id_to_gps.split("_")[:-4] + id_to_gps.split("_")[-3:])
+            )
+        else:
+            self.id_to_gps = torch.load(id_to_gps)
+        #self.unorm = UnormGPS()
+
+    def forward(self, x):
+        """Mapping from country id to gps coordinates
+        Args:
+            x: torch.Tensor with features
+        """
+
+        if isinstance(x, dict):
+            # for oracle
+            labels, x = x["label"], x["img"]
+        else:
+            # predicted labels
+            labels = x
+        self.id_to_gps = self.id_to_gps.to(labels.device)
+        #return {"gps": self.unorm(self.id_to_gps[labels])}
+        return {"gps": self.id_to_gps[labels]}

models/networks/heads/random.py

@@ -0,0 +1,53 @@
+import pandas as pd
+import torch
+from torch import nn
+from models.networks.utils import UnormGPS
+
+
+class Random(nn.Module):
+    def __init__(self, num_output):
+        """Random"""
+        super().__init__()
+        self.num_output = num_output
+        self.unorm = UnormGPS()
+
+    def forward(self, x):
+        """Predicts GPS coordinates from an image.
+        Args:
+            x: torch.Tensor with features
+        """
+        #x = x["img"]
+        gps = torch.rand((x.shape[0], self.num_output), device=x.device) * 2 - 1
+        return {"gps": self.unorm(gps)}
+
+
+class RandomCoords(nn.Module):
+    def __init__(self, coords_path: str):
+        """Randomly sample from a list of coordinates
+        Args:
+            coords_path: str with path to csv file with coordinates
+        """
+        super().__init__()
+        coordinates = pd.read_csv(coords_path)
+        longitudes = coordinates["longitude"].values / 180
+        latitudes = coordinates["latitude"].values / 90
+        self.unorm = UnormGPS()
+        del coordinates
+
+        self.N = len(longitudes)
+        assert len(longitudes) == len(latitudes)
+        self.coordinates = torch.stack(
+            [torch.tensor(latitudes), torch.tensor(longitudes)],
+            dim=-1,
+        )
+        del longitudes, latitudes
+
+    def forward(self, x):
+        """Predicts GPS coordinates from an image.
+        Args:
+            x: torch.Tensor with features
+        """
+        x = x["img"]
+        # randomly select a coordinate in the list
+        n = torch.randint(0, self.N, (x.shape[0],))
+        return {"gps": self.unorm(self.coordinates[n].to(x.device))}

models/networks/heads/regression.py

@@ -0,0 +1,44 @@
+from models.networks.utils import UnormGPS
+import torch.nn as nn
+from torch.nn.functional import tanh
+import torch
+
+
+class RegressionHead(nn.Module):
+    def __init__(self, use_tanh=False):
+        super().__init__()
+        self.unorm = UnormGPS()
+        self.use_tanh = use_tanh
+
+    def forward(self, x):
+        """Forward pass of the network.
+        x : Union[torch.Tensor, dict] with the output of the backbone.
+        """
+        if self.use_tanh:
+            x = tanh(x)
+        gps = self.unorm(x)
+        return {"gps": gps}
+
+
+class RegressionHeadAngle(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.unorm = UnormGPS()
+
+    def forward(self, x):
+        """Forward pass of the network.
+        x : Union[torch.Tensor, dict] with the output of the backbone.
+        """
+        x1 = x[:, 0].pow(2)
+        x2 = x[:, 1].pow(2)
+        x3 = x[:, 2].pow(2)
+        x4 = x[:, 3].pow(2)
+        cos_lambda = x1 / (x1 + x2)
+        sin_lambda = x2 / (x1 + x2)
+        cos_phi = x3 / (x3 + x4)
+        sin_phi = x4 / (x3 + x4)
+        lbd = torch.atan2(sin_lambda, cos_lambda)
+        phi = torch.atan2(sin_phi, cos_phi)
+        gps = torch.cat((lbd.unsqueeze(1), phi.unsqueeze(1)), dim=1)
+        # gps = self.unorm(x)
+        return {"gps": gps}