bug fix in localization

density_estimation.py

@@ -0,0 +1,80 @@
+import torch
+from torch.distributions import MultivariateNormal, Normal
+from torch.distributions.distribution import Distribution
+
+# Code from https://discuss.pytorch.org/t/kernel-density-estimation-as-loss-function/62261/8
+class GaussianKDE(Distribution):
+    def __init__(self, X, bw):
+        """
+        X : tensor (n, d)
+          `n` points with `d` dimensions to which KDE will be fit
+        bw : numeric
+          bandwidth for Gaussian kernel
+        """
+        self.X = X
+        self.bw = bw
+        self.dims = X.shape[-1]
+        self.n = X.shape[0]
+        self.mvn = MultivariateNormal(loc=torch.zeros(self.dims),
+                                      covariance_matrix=torch.eye(self.dims))
+
+    def sample(self, num_samples):
+        idxs = (np.random.uniform(0, 1, num_samples) * self.n).astype(int)
+        norm = Normal(loc=self.X[idxs], scale=self.bw)
+        return norm.sample()
+
+    def score_samples(self, Y, X=None):
+        """Returns the kernel density estimates of each point in `Y`.
+
+        Parameters
+        ----------
+        Y : tensor (m, d)
+          `m` points with `d` dimensions for which the probability density will
+          be calculated
+        X : tensor (n, d), optional
+          `n` points with `d` dimensions to which KDE will be fit. Provided to
+          allow batch calculations in `log_prob`. By default, `X` is None and
+          all points used to initialize KernelDensityEstimator are included.
+
+
+        Returns
+        -------
+        log_probs : tensor (m)
+          log probability densities for each of the queried points in `Y`
+        """
+        if X == None:
+            X = self.X
+        log_probs = torch.log(
+            (self.bw**(-self.dims) *
+             torch.exp(self.mvn.log_prob(
+                 (X.unsqueeze(1) - Y) / self.bw))).sum(dim=0) / self.n)
+
+        return log_probs
+
+    def log_prob(self, Y):
+        """Returns the total log probability of one or more points, `Y`, using
+        a Multivariate Normal kernel fit to `X` and scaled using `bw`.
+
+        Parameters
+        ----------
+        Y : tensor (m, d)
+          `m` points with `d` dimensions for which the probability density will
+          be calculated
+
+        Returns
+        -------
+        log_prob : numeric
+          total log probability density for the queried points, `Y`
+        """
+
+        X_chunks = self.X.split(1000)
+        Y_chunks = Y.split(1000)
+
+        log_prob = 0
+
+        for x in X_chunks:
+            for y in Y_chunks:
+                log_prob += self.score_samples(y, x).sum(dim=0)
+
+        return log_prob
+

localization.py

@@ -12,6 +12,7 @@
 from scipy import signal
 import torchvision.transforms as transforms
 import cv2
+import shutil
 
 class Localize:
     def __init__(self, model_weights, kernel_dim = (32,32), stride = 4, batch_size = 128, device = 'cuda', image_size = (256,256)):
@@ -132,12 +133,16 @@ def heatmap_on_image(image, hmap):
 
 def save_anomaly_map(image, hmap, save_path):
     imposed_image = heatmap_on_image(image, hmap)
-    cv2.imwrite(os.path.join(save_path, f'{file_name}.jpg'), image)
-    cv2.imwrite(os.path.join(self.sample_path, f'{file_name}_amap.jpg'), imposed_image)
+    file_name = image.split('/')[-1].split('.')[0]
+    shutil.copy(image, os.path.join(save_path, f'{file_name}.jpg'))
+    cv2.imwrite(os.path.join(save_path, f'{file_name}_amap.jpg'), imposed_image)
+
+
 
 L = Localize('./weights-bottle.ckpt')
 sp = L.patch_scores('./bottle/train/', './bottle/test/broken_large/004.png')
 GS = Gaussian_smoothing()
 up = GS.upsample(sp)
 visualize_heatmap('./bottle/test/broken_large/004.png', up)
 
+

train.py

@@ -42,8 +42,8 @@ def test_dataloader(self):
 
 
     def forward(self, x):
-        features, logits, embeds = self.model(x)
-        return features, logits, embeds
+        logits, embeds = self.model(x)
+        return  logits, embeds
 
     def configure_optimizers(self):
         optimizer = optim.SGD(self.parameters(), lr=self.hparams.learning_rate, 
@@ -61,7 +61,7 @@ def training_step(self, batch, batch_idx):
         x = torch.cat(batch, axis=0)
         y = torch.arange(len(batch))
         y = y.repeat_interleave(len(batch[0])).cuda()
-        features, logits, embeds = self(x)
+        logits, embeds = self(x)
         loss = self.criterion(logits, y)
         predicted = torch.argmax(logits,axis=1)
         accuracy = torch.true_divide(torch.sum(predicted==y), predicted.size(0))
@@ -72,7 +72,7 @@ def training_step(self, batch, batch_idx):
 
     def test_step(self, batch, batch_idx):
         x, y = batch
-        features, logits, embeds = self(x)
+        logits, embeds = self(x)
         loss = self.criterion(logits, y)