├── capture
    └── capture.jpg
├── requirements.txt
├── embeddings
    └── carpet
    │   └── embedding.pickle
├── sampling_methods
    ├── sampling_def.py
    └── kcenter_greedy.py
├── .gitignore
├── README.md
├── LICENSE
└── train.py


/capture/capture.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hcw-00/PatchCore_anomaly_detection/HEAD/capture/capture.jpg


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | opencv-python==4.5.2.52
2 | scikit-learn==0.24.2
3 | pytorch-lightning==1.3.3
4 | faiss-gpu==1.7.1.post3


--------------------------------------------------------------------------------
/embeddings/carpet/embedding.pickle:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hcw-00/PatchCore_anomaly_detection/HEAD/embeddings/carpet/embedding.pickle


--------------------------------------------------------------------------------
/sampling_methods/sampling_def.py:
--------------------------------------------------------------------------------
 1 | # Copyright 2017 Google Inc.
 2 | #
 3 | # Licensed under the Apache License, Version 2.0 (the "License");
 4 | # you may not use this file except in compliance with the License.
 5 | # You may obtain a copy of the License at
 6 | #
 7 | #      http://www.apache.org/licenses/LICENSE-2.0
 8 | #
 9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 | 
15 | """Abstract class for sampling methods.
16 | 
17 | Provides interface to sampling methods that allow same signature
18 | for select_batch.  Each subclass implements select_batch_ with the desired
19 | signature for readability.
20 | """
21 | 
22 | from __future__ import absolute_import
23 | from __future__ import division
24 | from __future__ import print_function
25 | 
26 | import abc
27 | import numpy as np
28 | 
29 | class SamplingMethod(object):
30 |   __metaclass__ = abc.ABCMeta
31 | 
32 |   @abc.abstractmethod
33 |   def __init__(self, X, y, seed, **kwargs):
34 |     self.X = X
35 |     self.y = y
36 |     self.seed = seed
37 | 
38 |   def flatten_X(self):
39 |     shape = self.X.shape
40 |     flat_X = self.X
41 |     if len(shape) > 2:
42 |       flat_X = np.reshape(self.X, (shape[0],np.product(shape[1:])))
43 |     return flat_X
44 | 
45 | 
46 |   @abc.abstractmethod
47 |   def select_batch_(self):
48 |     return
49 | 
50 |   def select_batch(self, **kwargs):
51 |     return self.select_batch_(**kwargs)
52 | 
53 |   def to_dict(self):
54 |     return None


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # Byte-compiled / optimized / DLL files
  2 | __pycache__/
  3 | *.py[cod]
  4 | *$py.class
  5 | 
  6 | # C extensions
  7 | *.so
  8 | 
  9 | # Distribution / packaging
 10 | .Python
 11 | build/
 12 | develop-eggs/
 13 | dist/
 14 | downloads/
 15 | eggs/
 16 | .eggs/
 17 | lib/
 18 | lib64/
 19 | parts/
 20 | sdist/
 21 | var/
 22 | wheels/
 23 | pip-wheel-metadata/
 24 | share/python-wheels/
 25 | *.egg-info/
 26 | .installed.cfg
 27 | *.egg
 28 | MANIFEST
 29 | 
 30 | # PyInstaller
 31 | #  Usually these files are written by a python script from a template
 32 | #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 33 | *.manifest
 34 | *.spec
 35 | 
 36 | # Installer logs
 37 | pip-log.txt
 38 | pip-delete-this-directory.txt
 39 | 
 40 | # Unit test / coverage reports
 41 | htmlcov/
 42 | .tox/
 43 | .nox/
 44 | .coverage
 45 | .coverage.*
 46 | .cache
 47 | nosetests.xml
 48 | coverage.xml
 49 | *.cover
 50 | *.py,cover
 51 | .hypothesis/
 52 | .pytest_cache/
 53 | 
 54 | # Translations
 55 | *.mo
 56 | *.pot
 57 | 
 58 | # Django stuff:
 59 | *.log
 60 | local_settings.py
 61 | db.sqlite3
 62 | db.sqlite3-journal
 63 | 
 64 | # Flask stuff:
 65 | instance/
 66 | .webassets-cache
 67 | 
 68 | # Scrapy stuff:
 69 | .scrapy
 70 | 
 71 | # Sphinx documentation
 72 | docs/_build/
 73 | 
 74 | # PyBuilder
 75 | target/
 76 | 
 77 | # Jupyter Notebook
 78 | .ipynb_checkpoints
 79 | 
 80 | # IPython
 81 | profile_default/
 82 | ipython_config.py
 83 | 
 84 | # pyenv
 85 | .python-version
 86 | 
 87 | # pipenv
 88 | #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
 89 | #   However, in case of collaboration, if having platform-specific dependencies or dependencies
 90 | #   having no cross-platform support, pipenv may install dependencies that don't work, or not
 91 | #   install all needed dependencies.
 92 | #Pipfile.lock
 93 | 
 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow
 95 | __pypackages__/
 96 | 
 97 | # Celery stuff
 98 | celerybeat-schedule
 99 | celerybeat.pid
100 | 
101 | # SageMath parsed files
102 | *.sage.py
103 | 
104 | # Environments
105 | .env
106 | .venv
107 | env/
108 | venv/
109 | ENV/
110 | env.bak/
111 | venv.bak/
112 | 
113 | # Spyder project settings
114 | .spyderproject
115 | .spyproject
116 | 
117 | # Rope project settings
118 | .ropeproject
119 | 
120 | # mkdocs documentation
121 | /site
122 | 
123 | # mypy
124 | .mypy_cache/
125 | .dmypy.json
126 | dmypy.json
127 | 
128 | # Pyre type checker
129 | .pyre/
130 | 
131 | # embedding.pickle
132 | ./embeddings/
133 | 
134 | *.png
135 | *.jpg
136 | *.ckpt
137 | test/*
138 | MVTec/*


--------------------------------------------------------------------------------
/sampling_methods/kcenter_greedy.py:
--------------------------------------------------------------------------------
  1 | # Copyright 2017 Google Inc.
  2 | #
  3 | # Licensed under the Apache License, Version 2.0 (the "License");
  4 | # you may not use this file except in compliance with the License.
  5 | # You may obtain a copy of the License at
  6 | #
  7 | #      http://www.apache.org/licenses/LICENSE-2.0
  8 | #
  9 | # Unless required by applicable law or agreed to in writing, software
 10 | # distributed under the License is distributed on an "AS IS" BASIS,
 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12 | # See the License for the specific language governing permissions and
 13 | # limitations under the License.
 14 | 
 15 | """Returns points that minimizes the maximum distance of any point to a center.
 16 | 
 17 | Implements the k-Center-Greedy method in
 18 | Ozan Sener and Silvio Savarese.  A Geometric Approach to Active Learning for
 19 | Convolutional Neural Networks. https://arxiv.org/abs/1708.00489 2017
 20 | 
 21 | Distance metric defaults to l2 distance.  Features used to calculate distance
 22 | are either raw features or if a model has transform method then uses the output
 23 | of model.transform(X).
 24 | 
 25 | Can be extended to a robust k centers algorithm that ignores a certain number of
 26 | outlier datapoints.  Resulting centers are solution to multiple integer program.
 27 | """
 28 | 
 29 | from __future__ import absolute_import
 30 | from __future__ import division
 31 | from __future__ import print_function
 32 | 
 33 | import numpy as np
 34 | from sklearn.metrics import pairwise_distances
 35 | from sampling_methods.sampling_def import SamplingMethod
 36 | 
 37 | 
 38 | class kCenterGreedy(SamplingMethod):
 39 | 
 40 |   def __init__(self, X, y, seed, metric='euclidean'):
 41 |     self.X = X
 42 |     self.y = y
 43 |     self.flat_X = self.flatten_X()
 44 |     self.name = 'kcenter'
 45 |     self.features = self.flat_X
 46 |     self.metric = metric
 47 |     self.min_distances = None
 48 |     self.n_obs = self.X.shape[0]
 49 |     self.already_selected = []
 50 | 
 51 |   def update_distances(self, cluster_centers, only_new=True, reset_dist=False):
 52 |     """Update min distances given cluster centers.
 53 | 
 54 |     Args:
 55 |       cluster_centers: indices of cluster centers
 56 |       only_new: only calculate distance for newly selected points and update
 57 |         min_distances.
 58 |       rest_dist: whether to reset min_distances.
 59 |     """
 60 | 
 61 |     if reset_dist:
 62 |       self.min_distances = None
 63 |     if only_new:
 64 |       cluster_centers = [d for d in cluster_centers
 65 |                          if d not in self.already_selected]
 66 |     if cluster_centers:
 67 |       # Update min_distances for all examples given new cluster center.
 68 |       x = self.features[cluster_centers]
 69 |       dist = pairwise_distances(self.features, x, metric=self.metric)
 70 | 
 71 |       if self.min_distances is None:
 72 |         self.min_distances = np.min(dist, axis=1).reshape(-1,1)
 73 |       else:
 74 |         self.min_distances = np.minimum(self.min_distances, dist)
 75 | 
 76 |   def select_batch_(self, model, already_selected, N, **kwargs):
 77 |     """
 78 |     Diversity promoting active learning method that greedily forms a batch
 79 |     to minimize the maximum distance to a cluster center among all unlabeled
 80 |     datapoints.
 81 | 
 82 |     Args:
 83 |       model: model with scikit-like API with decision_function implemented
 84 |       already_selected: index of datapoints already selected
 85 |       N: batch size
 86 | 
 87 |     Returns:
 88 |       indices of points selected to minimize distance to cluster centers
 89 |     """
 90 | 
 91 |     try:
 92 |       # Assumes that the transform function takes in original data and not
 93 |       # flattened data.
 94 |       print('Getting transformed features...')
 95 |       self.features = model.transform(self.X)
 96 |       print('Calculating distances...')
 97 |       self.update_distances(already_selected, only_new=False, reset_dist=True)
 98 |     except:
 99 |       print('Using flat_X as features.')
100 |       self.update_distances(already_selected, only_new=True, reset_dist=False)
101 | 
102 |     new_batch = []
103 | 
104 |     for _ in range(N):
105 |       if self.already_selected is None:
106 |         # Initialize centers with a randomly selected datapoint
107 |         ind = np.random.choice(np.arange(self.n_obs))
108 |       else:
109 |         ind = np.argmax(self.min_distances)
110 |       # New examples should not be in already selected since those points
111 |       # should have min_distance of zero to a cluster center.
112 |       assert ind not in already_selected
113 | 
114 |       self.update_distances([ind], only_new=True, reset_dist=False)
115 |       new_batch.append(ind)
116 |     print('Maximum distance from cluster centers is %0.2f'
117 |             % max(self.min_distances))
118 | 
119 | 
120 |     self.already_selected = already_selected
121 | 
122 |     return new_batch
123 | 
124 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # PatchCore anomaly detection
  2 | Unofficial implementation of PatchCore(new SOTA) anomaly detection model
  3 | 
  4 | 
  5 | Original Paper : 
  6 | Towards Total Recall in Industrial Anomaly Detection (Jun 2021)  
  7 | Karsten Roth, Latha Pemula, Joaquin Zepeda, Bernhard Schölkopf, Thomas Brox, Peter Gehler  
  8 | 
  9 | 
 10 | https://arxiv.org/abs/2106.08265  
 11 | https://paperswithcode.com/sota/anomaly-detection-on-mvtec-ad
 12 | 
 13 | ![plot](./capture/capture.jpg)
 14 | 
 15 | 
 16 | updates(21/06/21) :  
 17 | - I used sklearn's SparseRandomProjection(ep=0.9) for random projection. I'm not confident with this.
 18 | - I think exact value of "b nearest patch-features" is not presented in the paper. I just set 9. (args.n_neighbors)  
 19 | - In terms of NN search, author used "faiss". but not implemented in this code yet. 
 20 | - sample embeddings/carpet/embedding.pickle => coreset_sampling_ratio=0.001  
 21 | 
 22 | updates(21/06/26) :  
 23 | - A critical [issue](https://github.com/hcw-00/PatchCore_anomaly_detection/issues/3#issue-930229038) related to "locally aware patch" raised and fixed. Score table is updated. 
 24 | 
 25 | ### Usage 
 26 | ~~~
 27 | # install python 3.6, torch==1.8.1, torchvision==0.9.1
 28 | pip install -r requirements.txt
 29 | 
 30 | python train.py --phase train or test --dataset_path .../mvtec_anomaly_detection --category carpet --project_root_path path/to/save/results --coreset_sampling_ratio 0.01 --n_neighbors 9'
 31 | 
 32 | # for fast try just specify your dataset_path and run
 33 | python train.py --phase test --dataset_path .../mvtec_anomaly_detection --project_root_path ./
 34 | ~~~
 35 | 
 36 | ### MVTecAD AUROC score (PatchCore-1%, mean of n trials)
 37 | | Category | Paper<br>(image-level) | This code<br>(image-level) | Paper<br>(pixel-level) | This code<br>(pixel-level) |
 38 | | :-----: | :-: | :-: | :-: | :-: |
 39 | | carpet | 0.980 | 0.991(1) | 0.989 | 0.989(1) |
 40 | | grid | 0.986 | 0.975(1) | 0.986 | 0.975(1) |
 41 | | leather | 1.000 | 1.000(1) | 0.993 | 0.991(1) |
 42 | | tile | 0.994 | 0.994(1) | 0.961 | 0.949(1) |
 43 | | wood | 0.992 | 0.989(1) | 0.951 | 0.936(1) |
 44 | | bottle | 1.000 | 1.000(1) | 0.985 | 0.981(1) |
 45 | | cable | 0.993 | 0.995(1) | 0.982 | 0.983(1) |
 46 | | capsule | 0.980 | 0.976(1) | 0.988 | 0.989(1) |
 47 | | hazelnut | 1.000 | 1.000(1) | 0.986 | 0.985(1) |
 48 | | metal nut | 0.997 | 0.999(1) | 0.984 | 0.984(1) |
 49 | | pill | 0.970 | 0.959(1) | 0.971 | 0.977(1) |
 50 | | screw | 0.964 | 0.949(1) | 0.992 | 0.977(1) |
 51 | | toothbrush | 1.000 | 1.000(1) | 0.985 | 0.986(1) |
 52 | | transistor | 0.999 | 1.000(1) | 0.949 | 0.972(1) |
 53 | | zipper | 0.992 | 0.995(1) | 0.988 | 0.984(1) |
 54 | | mean | 0.990 | 0.988 | 0.980 | 0.977 |
 55 | 
 56 | ### Code Reference
 57 | kcenter algorithm :  
 58 | https://github.com/google/active-learning  
 59 | embedding concat function :  
 60 | https://github.com/xiahaifeng1995/PaDiM-Anomaly-Detection-Localization-master
 61 | 
 62 | <!--
 63 | ### MVTecAD AUROC score (PatchCore-1%, mean of n trials)
 64 | | Category | Paper<br>(image-level) | This code<br>(image-level) | Paper<br>(pixel-level) | This code<br>(pixel-level) |
 65 | | :-----: | :-: | :-: | :-: | :-: |
 66 | | carpet | 0.980 | 0.997(1) | 0.989 | 0.990(1) |
 67 | | grid | 0.986 | 0.941(1) | 0.986 | 0.983(1) |
 68 | | leather | 1.000 | 1.000(1) | 0.993 | 0.991(1) |
 69 | | tile | 0.994 | 0.982(1) | 0.961 | 0.932(1) |
 70 | | wood | 0.992 | 0.999(1) | 0.951 | 0.976(1) |
 71 | | bottle | 1.000 | 0.986(1) | 0.985 | 0.941(1) |
 72 | | cable | 0.993 | 0.970(1) | 0.982 | 0.955(1) |
 73 | | capsule | 0.980 | 0.949(1) | 0.988 | 0.987(1) |
 74 | | hazelnut | 1.000 | 0.997(1) | 0.986 | 0.982(1) |
 75 | | metal nut | 0.997 | 0.997(1) | 0.984 | 0.962(1) |
 76 | | pill | 0.970 | 0.918(1) | 0.971 | 0.941(1) |
 77 | | screw | 0.964 | 0.967(1) | 0.992 | 0.987(1) |
 78 | | toothbrush | 1.000 | 0.997(1) | 0.985 | 0.984(1) |
 79 | | transistor | 0.999 | 0.960(1) | 0.949 | 0.894(1) |
 80 | | zipper | 0.992 | 0.968(1) | 0.988 | 0.987(1) |
 81 | | mean | 0.990 | 0.975 | 0.980 | 0.966 |
 82 | -->
 83 | <!--
 84 | carpet
 85 | {'img_auc': 0.9907704654895666, 'pixel_auc': 0.988687705858456}
 86 | grid
 87 | {'img_auc': 0.974937343358396, 'pixel_auc': 0.9751171484318466}
 88 | leather
 89 | {'img_auc': 1.0, 'pixel_auc': 0.990535126453962}
 90 | tile
 91 | {'img_auc': 0.9935064935064934, 'pixel_auc': 0.9487486588602011}
 92 | wood
 93 | {'img_auc': 0.9894736842105263, 'pixel_auc': 0.9361264985782104}
 94 | bottle
 95 | {'img_auc': 1.0, 'pixel_auc': 0.9812954295492088}
 96 | cable
 97 | {'img_auc': 0.994752623688156, 'pixel_auc': 0.9827989792867079}
 98 | capsule
 99 | {'img_auc': 0.9764658954926206, 'pixel_auc': 0.9889956357045827}
100 | hazelnut
101 | {'img_auc': 1.0, 'pixel_auc': 0.9846421467788966}
102 | matal nut
103 | {'img_auc': 0.9990224828934506, 'pixel_auc': 0.9838413598506198}
104 | pill
105 | {'img_auc': 0.9593562465902891, 'pixel_auc': 0.9772725707677767}
106 | screw
107 | {'img_auc': 0.9491699118671859, 'pixel_auc': 0.9774047713027145}
108 | toothbrush
109 | {'img_auc': 1.0, 'pixel_auc': 0.9862281871506898}
110 | transistor
111 | {'img_auc': 0.9995833333333333, 'pixel_auc': 0.9715171862282481}
112 | zipper
113 | {'img_auc': 0.9952731092436975, 'pixel_auc': 0.9838086200703351}
114 | >
115 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
  1 |                                  Apache License
  2 |                            Version 2.0, January 2004
  3 |                         http://www.apache.org/licenses/
  4 | 
  5 |    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
  6 | 
  7 |    1. Definitions.
  8 | 
  9 |       "License" shall mean the terms and conditions for use, reproduction,
 10 |       and distribution as defined by Sections 1 through 9 of this document.
 11 | 
 12 |       "Licensor" shall mean the copyright owner or entity authorized by
 13 |       the copyright owner that is granting the License.
 14 | 
 15 |       "Legal Entity" shall mean the union of the acting entity and all
 16 |       other entities that control, are controlled by, or are under common
 17 |       control with that entity. For the purposes of this definition,
 18 |       "control" means (i) the power, direct or indirect, to cause the
 19 |       direction or management of such entity, whether by contract or
 20 |       otherwise, or (ii) ownership of fifty percent (50%) or more of the
 21 |       outstanding shares, or (iii) beneficial ownership of such entity.
 22 | 
 23 |       "You" (or "Your") shall mean an individual or Legal Entity
 24 |       exercising permissions granted by this License.
 25 | 
 26 |       "Source" form shall mean the preferred form for making modifications,
 27 |       including but not limited to software source code, documentation
 28 |       source, and configuration files.
 29 | 
 30 |       "Object" form shall mean any form resulting from mechanical
 31 |       transformation or translation of a Source form, including but
 32 |       not limited to compiled object code, generated documentation,
 33 |       and conversions to other media types.
 34 | 
 35 |       "Work" shall mean the work of authorship, whether in Source or
 36 |       Object form, made available under the License, as indicated by a
 37 |       copyright notice that is included in or attached to the work
 38 |       (an example is provided in the Appendix below).
 39 | 
 40 |       "Derivative Works" shall mean any work, whether in Source or Object
 41 |       form, that is based on (or derived from) the Work and for which the
 42 |       editorial revisions, annotations, elaborations, or other modifications
 43 |       represent, as a whole, an original work of authorship. For the purposes
 44 |       of this License, Derivative Works shall not include works that remain
 45 |       separable from, or merely link (or bind by name) to the interfaces of,
 46 |       the Work and Derivative Works thereof.
 47 | 
 48 |       "Contribution" shall mean any work of authorship, including
 49 |       the original version of the Work and any modifications or additions
 50 |       to that Work or Derivative Works thereof, that is intentionally
 51 |       submitted to Licensor for inclusion in the Work by the copyright owner
 52 |       or by an individual or Legal Entity authorized to submit on behalf of
 53 |       the copyright owner. For the purposes of this definition, "submitted"
 54 |       means any form of electronic, verbal, or written communication sent
 55 |       to the Licensor or its representatives, including but not limited to
 56 |       communication on electronic mailing lists, source code control systems,
 57 |       and issue tracking systems that are managed by, or on behalf of, the
 58 |       Licensor for the purpose of discussing and improving the Work, but
 59 |       excluding communication that is conspicuously marked or otherwise
 60 |       designated in writing by the copyright owner as "Not a Contribution."
 61 | 
 62 |       "Contributor" shall mean Licensor and any individual or Legal Entity
 63 |       on behalf of whom a Contribution has been received by Licensor and
 64 |       subsequently incorporated within the Work.
 65 | 
 66 |    2. Grant of Copyright License. Subject to the terms and conditions of
 67 |       this License, each Contributor hereby grants to You a perpetual,
 68 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 69 |       copyright license to reproduce, prepare Derivative Works of,
 70 |       publicly display, publicly perform, sublicense, and distribute the
 71 |       Work and such Derivative Works in Source or Object form.
 72 | 
 73 |    3. Grant of Patent License. Subject to the terms and conditions of
 74 |       this License, each Contributor hereby grants to You a perpetual,
 75 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 76 |       (except as stated in this section) patent license to make, have made,
 77 |       use, offer to sell, sell, import, and otherwise transfer the Work,
 78 |       where such license applies only to those patent claims licensable
 79 |       by such Contributor that are necessarily infringed by their
 80 |       Contribution(s) alone or by combination of their Contribution(s)
 81 |       with the Work to which such Contribution(s) was submitted. If You
 82 |       institute patent litigation against any entity (including a
 83 |       cross-claim or counterclaim in a lawsuit) alleging that the Work
 84 |       or a Contribution incorporated within the Work constitutes direct
 85 |       or contributory patent infringement, then any patent licenses
 86 |       granted to You under this License for that Work shall terminate
 87 |       as of the date such litigation is filed.
 88 | 
 89 |    4. Redistribution. You may reproduce and distribute copies of the
 90 |       Work or Derivative Works thereof in any medium, with or without
 91 |       modifications, and in Source or Object form, provided that You
 92 |       meet the following conditions:
 93 | 
 94 |       (a) You must give any other recipients of the Work or
 95 |           Derivative Works a copy of this License; and
 96 | 
 97 |       (b) You must cause any modified files to carry prominent notices
 98 |           stating that You changed the files; and
 99 | 
100 |       (c) You must retain, in the Source form of any Derivative Works
101 |           that You distribute, all copyright, patent, trademark, and
102 |           attribution notices from the Source form of the Work,
103 |           excluding those notices that do not pertain to any part of
104 |           the Derivative Works; and
105 | 
106 |       (d) If the Work includes a "NOTICE" text file as part of its
107 |           distribution, then any Derivative Works that You distribute must
108 |           include a readable copy of the attribution notices contained
109 |           within such NOTICE file, excluding those notices that do not
110 |           pertain to any part of the Derivative Works, in at least one
111 |           of the following places: within a NOTICE text file distributed
112 |           as part of the Derivative Works; within the Source form or
113 |           documentation, if provided along with the Derivative Works; or,
114 |           within a display generated by the Derivative Works, if and
115 |           wherever such third-party notices normally appear. The contents
116 |           of the NOTICE file are for informational purposes only and
117 |           do not modify the License. You may add Your own attribution
118 |           notices within Derivative Works that You distribute, alongside
119 |           or as an addendum to the NOTICE text from the Work, provided
120 |           that such additional attribution notices cannot be construed
121 |           as modifying the License.
122 | 
123 |       You may add Your own copyright statement to Your modifications and
124 |       may provide additional or different license terms and conditions
125 |       for use, reproduction, or distribution of Your modifications, or
126 |       for any such Derivative Works as a whole, provided Your use,
127 |       reproduction, and distribution of the Work otherwise complies with
128 |       the conditions stated in this License.
129 | 
130 |    5. Submission of Contributions. Unless You explicitly state otherwise,
131 |       any Contribution intentionally submitted for inclusion in the Work
132 |       by You to the Licensor shall be under the terms and conditions of
133 |       this License, without any additional terms or conditions.
134 |       Notwithstanding the above, nothing herein shall supersede or modify
135 |       the terms of any separate license agreement you may have executed
136 |       with Licensor regarding such Contributions.
137 | 
138 |    6. Trademarks. This License does not grant permission to use the trade
139 |       names, trademarks, service marks, or product names of the Licensor,
140 |       except as required for reasonable and customary use in describing the
141 |       origin of the Work and reproducing the content of the NOTICE file.
142 | 
143 |    7. Disclaimer of Warranty. Unless required by applicable law or
144 |       agreed to in writing, Licensor provides the Work (and each
145 |       Contributor provides its Contributions) on an "AS IS" BASIS,
146 |       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 |       implied, including, without limitation, any warranties or conditions
148 |       of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 |       PARTICULAR PURPOSE. You are solely responsible for determining the
150 |       appropriateness of using or redistributing the Work and assume any
151 |       risks associated with Your exercise of permissions under this License.
152 | 
153 |    8. Limitation of Liability. In no event and under no legal theory,
154 |       whether in tort (including negligence), contract, or otherwise,
155 |       unless required by applicable law (such as deliberate and grossly
156 |       negligent acts) or agreed to in writing, shall any Contributor be
157 |       liable to You for damages, including any direct, indirect, special,
158 |       incidental, or consequential damages of any character arising as a
159 |       result of this License or out of the use or inability to use the
160 |       Work (including but not limited to damages for loss of goodwill,
161 |       work stoppage, computer failure or malfunction, or any and all
162 |       other commercial damages or losses), even if such Contributor
163 |       has been advised of the possibility of such damages.
164 | 
165 |    9. Accepting Warranty or Additional Liability. While redistributing
166 |       the Work or Derivative Works thereof, You may choose to offer,
167 |       and charge a fee for, acceptance of support, warranty, indemnity,
168 |       or other liability obligations and/or rights consistent with this
169 |       License. However, in accepting such obligations, You may act only
170 |       on Your own behalf and on Your sole responsibility, not on behalf
171 |       of any other Contributor, and only if You agree to indemnify,
172 |       defend, and hold each Contributor harmless for any liability
173 |       incurred by, or claims asserted against, such Contributor by reason
174 |       of your accepting any such warranty or additional liability.
175 | 
176 |    END OF TERMS AND CONDITIONS
177 | 
178 |    APPENDIX: How to apply the Apache License to your work.
179 | 
180 |       To apply the Apache License to your work, attach the following
181 |       boilerplate notice, with the fields enclosed by brackets "[]"
182 |       replaced with your own identifying information. (Don't include
183 |       the brackets!)  The text should be enclosed in the appropriate
184 |       comment syntax for the file format. We also recommend that a
185 |       file or class name and description of purpose be included on the
186 |       same "printed page" as the copyright notice for easier
187 |       identification within third-party archives.
188 | 
189 |    Copyright [yyyy] [name of copyright owner]
190 | 
191 |    Licensed under the Apache License, Version 2.0 (the "License");
192 |    you may not use this file except in compliance with the License.
193 |    You may obtain a copy of the License at
194 | 
195 |        http://www.apache.org/licenses/LICENSE-2.0
196 | 
197 |    Unless required by applicable law or agreed to in writing, software
198 |    distributed under the License is distributed on an "AS IS" BASIS,
199 |    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 |    See the License for the specific language governing permissions and
201 |    limitations under the License.
202 | 


--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import glob
  3 | import shutil
  4 | 
  5 | import numpy as np
  6 | from PIL import Image
  7 | import cv2
  8 | 
  9 | from sklearn.random_projection import SparseRandomProjection
 10 | from sklearn.metrics import roc_auc_score
 11 | from sampling_methods.kcenter_greedy import kCenterGreedy
 12 | from scipy.ndimage import gaussian_filter
 13 | import torch
 14 | from torch.nn import functional as F
 15 | from torchvision import transforms
 16 | from torch.utils.data import Dataset, DataLoader
 17 | import pytorch_lightning as pl
 18 | import faiss
 19 | 
 20 | 
 21 | def distance_matrix(x, y=None, p=2):  # pairwise distance of vectors
 22 |     y = x if type(y) == type(None) else y
 23 | 
 24 |     n = x.size(0)
 25 |     m = y.size(0)
 26 |     d = x.size(1)
 27 | 
 28 |     x = x.unsqueeze(1).expand(n, m, d)
 29 |     y = y.unsqueeze(0).expand(n, m, d)
 30 | 
 31 |     dist = torch.pow(x - y, p).sum(2)
 32 |     return dist
 33 | 
 34 | 
 35 | class NN():
 36 |     def __init__(self, X=None, Y=None, p=2):
 37 |         self.p = p
 38 |         self.train(X, Y)
 39 | 
 40 |     def train(self, X, Y):
 41 |         self.train_pts = X
 42 |         self.train_label = Y
 43 | 
 44 |     def __call__(self, x):
 45 |         return self.predict(x)
 46 | 
 47 |     def predict(self, x):
 48 |         if type(self.train_pts) == type(None) or type(self.train_label) == type(None):
 49 |             name = self.__class__.__name__
 50 |             raise RuntimeError(f"{name} wasn't trained. Need to execute {name}.train() first")
 51 | 
 52 |         dist = distance_matrix(x, self.train_pts, self.p) ** (1 / self.p)
 53 |         labels = torch.argmin(dist, dim=1)
 54 |         return self.train_label[labels]
 55 | 
 56 | class KNN(NN):
 57 |     def __init__(self, X=None, Y=None, k=3, p=2):
 58 |         self.k = k
 59 |         super().__init__(X, Y, p)
 60 | 
 61 |     def train(self, X, Y):
 62 |         super().train(X, Y)
 63 |         if type(Y) != type(None):
 64 |             self.unique_labels = self.train_label.unique()
 65 | 
 66 |     def predict(self, x):
 67 |         dist = torch.cdist(x, self.train_pts, self.p)
 68 |         knn = dist.topk(self.k, largest=False)
 69 |         return knn
 70 | 
 71 | def prep_dirs(root):
 72 |     # make embeddings dir
 73 |     embeddings_path = os.path.join('./', 'embeddings', args.category)
 74 |     os.makedirs(embeddings_path, exist_ok=True)
 75 |     # make sample dir
 76 |     sample_path = os.path.join(root, 'sample')
 77 |     os.makedirs(sample_path, exist_ok=True)
 78 |     # make source code record dir & copy
 79 |     source_code_save_path = os.path.join(root, 'src')
 80 |     os.makedirs(source_code_save_path, exist_ok=True)
 81 |     return embeddings_path, sample_path, source_code_save_path
 82 | 
 83 | def embedding_concat(x, y):
 84 |     # from https://github.com/xiahaifeng1995/PaDiM-Anomaly-Detection-Localization-master
 85 |     B, C1, H1, W1 = x.size()
 86 |     _, C2, H2, W2 = y.size()
 87 |     s = int(H1 / H2)
 88 |     x = F.unfold(x, kernel_size=s, dilation=1, stride=s)
 89 |     x = x.view(B, C1, -1, H2, W2)
 90 |     z = torch.zeros(B, C1 + C2, x.size(2), H2, W2)
 91 |     for i in range(x.size(2)):
 92 |         z[:, :, i, :, :] = torch.cat((x[:, :, i, :, :], y), 1)
 93 |     z = z.view(B, -1, H2 * W2)
 94 |     z = F.fold(z, kernel_size=s, output_size=(H1, W1), stride=s)
 95 |     return z
 96 | 
 97 | def reshape_embedding(embedding):
 98 |     embedding_list = []
 99 |     for k in range(embedding.shape[0]):
100 |         for i in range(embedding.shape[2]):
101 |             for j in range(embedding.shape[3]):
102 |                 embedding_list.append(embedding[k, :, i, j])
103 |     return embedding_list
104 | 
105 | #imagenet
106 | mean_train = [0.485, 0.456, 0.406]
107 | std_train = [0.229, 0.224, 0.225]
108 | 
109 | class MVTecDataset(Dataset):
110 |     def __init__(self, root, transform, gt_transform, phase):
111 |         if phase=='train':
112 |             self.img_path = os.path.join(root, 'train')
113 |         else:
114 |             self.img_path = os.path.join(root, 'test')
115 |             self.gt_path = os.path.join(root, 'ground_truth')
116 |         self.transform = transform
117 |         self.gt_transform = gt_transform
118 |         # load dataset
119 |         self.img_paths, self.gt_paths, self.labels, self.types = self.load_dataset() # self.labels => good : 0, anomaly : 1
120 | 
121 |     def load_dataset(self):
122 | 
123 |         img_tot_paths = []
124 |         gt_tot_paths = []
125 |         tot_labels = []
126 |         tot_types = []
127 | 
128 |         defect_types = os.listdir(self.img_path)
129 |         
130 |         for defect_type in defect_types:
131 |             if defect_type == 'good':
132 |                 img_paths = glob.glob(os.path.join(self.img_path, defect_type) + "/*.png")
133 |                 img_tot_paths.extend(img_paths)
134 |                 gt_tot_paths.extend([0]*len(img_paths))
135 |                 tot_labels.extend([0]*len(img_paths))
136 |                 tot_types.extend(['good']*len(img_paths))
137 |             else:
138 |                 img_paths = glob.glob(os.path.join(self.img_path, defect_type) + "/*.png")
139 |                 gt_paths = glob.glob(os.path.join(self.gt_path, defect_type) + "/*.png")
140 |                 img_paths.sort()
141 |                 gt_paths.sort()
142 |                 img_tot_paths.extend(img_paths)
143 |                 gt_tot_paths.extend(gt_paths)
144 |                 tot_labels.extend([1]*len(img_paths))
145 |                 tot_types.extend([defect_type]*len(img_paths))
146 | 
147 |         assert len(img_tot_paths) == len(gt_tot_paths), "Something wrong with test and ground truth pair!"
148 |         
149 |         return img_tot_paths, gt_tot_paths, tot_labels, tot_types
150 | 
151 |     def __len__(self):
152 |         return len(self.img_paths)
153 | 
154 |     def __getitem__(self, idx):
155 |         img_path, gt, label, img_type = self.img_paths[idx], self.gt_paths[idx], self.labels[idx], self.types[idx]
156 |         img = Image.open(img_path).convert('RGB')
157 |         img = self.transform(img)
158 |         if gt == 0:
159 |             gt = torch.zeros([1, img.size()[-2], img.size()[-2]])
160 |         else:
161 |             gt = Image.open(gt)
162 |             gt = self.gt_transform(gt)
163 |         
164 |         assert img.size()[1:] == gt.size()[1:], "image.size != gt.size !!!"
165 | 
166 |         return img, gt, label, os.path.basename(img_path[:-4]), img_type
167 | 
168 | def cvt2heatmap(gray):
169 |     heatmap = cv2.applyColorMap(np.uint8(gray), cv2.COLORMAP_JET)
170 |     return heatmap
171 | 
172 | def heatmap_on_image(heatmap, image):
173 |     if heatmap.shape != image.shape:
174 |         heatmap = cv2.resize(heatmap, (image.shape[0], image.shape[1]))
175 |     out = np.float32(heatmap)/255 + np.float32(image)/255
176 |     out = out / np.max(out)
177 |     return np.uint8(255 * out)
178 | 
179 | def min_max_norm(image):
180 |     a_min, a_max = image.min(), image.max()
181 |     return (image-a_min)/(a_max - a_min)    
182 |     
183 | 
184 | class PatchCore(pl.LightningModule):
185 |     def __init__(self, hparams):
186 |         super(PatchCore, self).__init__()
187 | 
188 |         self.save_hyperparameters(hparams)
189 | 
190 |         self.init_features()
191 |         def hook_t(module, input, output):
192 |             self.features.append(output)
193 | 
194 |         self.model = torch.hub.load('pytorch/vision:v0.9.0', 'wide_resnet50_2', pretrained=True)
195 | 
196 |         for param in self.model.parameters():
197 |             param.requires_grad = False
198 | 
199 |         self.model.layer2[-1].register_forward_hook(hook_t)
200 |         self.model.layer3[-1].register_forward_hook(hook_t)
201 | 
202 |         self.criterion = torch.nn.MSELoss(reduction='sum')
203 | 
204 |         self.init_results_list()
205 | 
206 |         self.data_transforms = transforms.Compose([
207 |                         transforms.Resize((args.load_size, args.load_size), Image.ANTIALIAS),
208 |                         transforms.ToTensor(),
209 |                         transforms.CenterCrop(args.input_size),
210 |                         transforms.Normalize(mean=mean_train,
211 |                                             std=std_train)])
212 |         self.gt_transforms = transforms.Compose([
213 |                         transforms.Resize((args.load_size, args.load_size)),
214 |                         transforms.ToTensor(),
215 |                         transforms.CenterCrop(args.input_size)])
216 | 
217 |         self.inv_normalize = transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.255], std=[1/0.229, 1/0.224, 1/0.255])
218 | 
219 |     def init_results_list(self):
220 |         self.gt_list_px_lvl = []
221 |         self.pred_list_px_lvl = []
222 |         self.gt_list_img_lvl = []
223 |         self.pred_list_img_lvl = []
224 |         self.img_path_list = []        
225 | 
226 |     def init_features(self):
227 |         self.features = []
228 | 
229 |     def forward(self, x_t):
230 |         self.init_features()
231 |         _ = self.model(x_t)
232 |         return self.features
233 | 
234 |     def save_anomaly_map(self, anomaly_map, input_img, gt_img, file_name, x_type):
235 |         if anomaly_map.shape != input_img.shape:
236 |             anomaly_map = cv2.resize(anomaly_map, (input_img.shape[0], input_img.shape[1]))
237 |         anomaly_map_norm = min_max_norm(anomaly_map)
238 |         anomaly_map_norm_hm = cvt2heatmap(anomaly_map_norm*255)
239 | 
240 |         # anomaly map on image
241 |         heatmap = cvt2heatmap(anomaly_map_norm*255)
242 |         hm_on_img = heatmap_on_image(heatmap, input_img)
243 | 
244 |         # save images
245 |         cv2.imwrite(os.path.join(self.sample_path, f'{x_type}_{file_name}.jpg'), input_img)
246 |         cv2.imwrite(os.path.join(self.sample_path, f'{x_type}_{file_name}_amap.jpg'), anomaly_map_norm_hm)
247 |         cv2.imwrite(os.path.join(self.sample_path, f'{x_type}_{file_name}_amap_on_img.jpg'), hm_on_img)
248 |         cv2.imwrite(os.path.join(self.sample_path, f'{x_type}_{file_name}_gt.jpg'), gt_img)
249 | 
250 |     def train_dataloader(self):
251 |         image_datasets = MVTecDataset(root=os.path.join(args.dataset_path,args.category), transform=self.data_transforms, gt_transform=self.gt_transforms, phase='train')
252 |         train_loader = DataLoader(image_datasets, batch_size=args.batch_size, shuffle=True, num_workers=0)
253 |         return train_loader
254 | 
255 |     def test_dataloader(self):
256 |         test_datasets = MVTecDataset(root=os.path.join(args.dataset_path,args.category), transform=self.data_transforms, gt_transform=self.gt_transforms, phase='test')
257 |         test_loader = DataLoader(test_datasets, batch_size=1, shuffle=False, num_workers=0)
258 |         return test_loader
259 | 
260 |     def configure_optimizers(self):
261 |         return None
262 | 
263 |     def on_train_start(self):
264 |         self.model.eval() # to stop running_var move (maybe not critical)        
265 |         self.embedding_dir_path, self.sample_path, self.source_code_save_path = prep_dirs(self.logger.log_dir)
266 |         self.embedding_list = []
267 |     
268 |     def on_test_start(self):
269 |         self.embedding_dir_path, self.sample_path, self.source_code_save_path = prep_dirs(self.logger.log_dir)
270 |         self.index = faiss.read_index(os.path.join(self.embedding_dir_path,'index.faiss'))
271 |         if torch.cuda.is_available():
272 |             res = faiss.StandardGpuResources()
273 |             self.index = faiss.index_cpu_to_gpu(res, 0 ,self.index)
274 |         self.init_results_list()
275 |         
276 |     def training_step(self, batch, batch_idx): # save locally aware patch features
277 |         x, _, _, _, _ = batch
278 |         features = self(x)
279 |         embeddings = []
280 |         for feature in features:
281 |             m = torch.nn.AvgPool2d(3, 1, 1)
282 |             embeddings.append(m(feature))
283 |         embedding = embedding_concat(embeddings[0], embeddings[1])
284 |         self.embedding_list.extend(reshape_embedding(np.array(embedding)))
285 | 
286 |     def training_epoch_end(self, outputs): 
287 |         total_embeddings = np.array(self.embedding_list)
288 |         # Random projection
289 |         self.randomprojector = SparseRandomProjection(n_components='auto', eps=0.9) # 'auto' => Johnson-Lindenstrauss lemma
290 |         self.randomprojector.fit(total_embeddings)
291 |         # Coreset Subsampling
292 |         selector = kCenterGreedy(total_embeddings,0,0)
293 |         selected_idx = selector.select_batch(model=self.randomprojector, already_selected=[], N=int(total_embeddings.shape[0]*args.coreset_sampling_ratio))
294 |         self.embedding_coreset = total_embeddings[selected_idx]
295 |         
296 |         print('initial embedding size : ', total_embeddings.shape)
297 |         print('final embedding size : ', self.embedding_coreset.shape)
298 |         #faiss
299 |         self.index = faiss.IndexFlatL2(self.embedding_coreset.shape[1])
300 |         self.index.add(self.embedding_coreset) 
301 |         faiss.write_index(self.index,  os.path.join(self.embedding_dir_path,'index.faiss'))
302 | 
303 | 
304 |     def test_step(self, batch, batch_idx): # Nearest Neighbour Search
305 |         x, gt, label, file_name, x_type = batch
306 |         # extract embedding
307 |         features = self(x)
308 |         embeddings = []
309 |         for feature in features:
310 |             m = torch.nn.AvgPool2d(3, 1, 1)
311 |             embeddings.append(m(feature))
312 |         embedding_ = embedding_concat(embeddings[0], embeddings[1])
313 |         embedding_test = np.array(reshape_embedding(np.array(embedding_)))
314 |         score_patches, _ = self.index.search(embedding_test , k=args.n_neighbors)
315 |         anomaly_map = score_patches[:,0].reshape((28,28))
316 |         N_b = score_patches[np.argmax(score_patches[:,0])]
317 |         w = (1 - (np.max(np.exp(N_b))/np.sum(np.exp(N_b))))
318 |         score = w*max(score_patches[:,0]) # Image-level score
319 |         gt_np = gt.cpu().numpy()[0,0].astype(int)
320 |         anomaly_map_resized = cv2.resize(anomaly_map, (args.input_size, args.input_size))
321 |         anomaly_map_resized_blur = gaussian_filter(anomaly_map_resized, sigma=4)
322 |         self.gt_list_px_lvl.extend(gt_np.ravel())
323 |         self.pred_list_px_lvl.extend(anomaly_map_resized_blur.ravel())
324 |         self.gt_list_img_lvl.append(label.cpu().numpy()[0])
325 |         self.pred_list_img_lvl.append(score)
326 |         self.img_path_list.extend(file_name)
327 |         # save images
328 |         x = self.inv_normalize(x)
329 |         input_x = cv2.cvtColor(x.permute(0,2,3,1).cpu().numpy()[0]*255, cv2.COLOR_BGR2RGB)
330 |         self.save_anomaly_map(anomaly_map_resized_blur, input_x, gt_np*255, file_name[0], x_type[0])
331 | 
332 |     def test_epoch_end(self, outputs):
333 |         print("Total pixel-level auc-roc score :")
334 |         pixel_auc = roc_auc_score(self.gt_list_px_lvl, self.pred_list_px_lvl)
335 |         print(pixel_auc)
336 |         print("Total image-level auc-roc score :")
337 |         img_auc = roc_auc_score(self.gt_list_img_lvl, self.pred_list_img_lvl)
338 |         print(img_auc)
339 |         print('test_epoch_end')
340 |         values = {'pixel_auc': pixel_auc, 'img_auc': img_auc}
341 |         self.log_dict(values)
342 | 
343 | def get_args():
344 |     import argparse
345 |     parser = argparse.ArgumentParser(description='ANOMALYDETECTION')
346 |     parser.add_argument('--phase', choices=['train','test'], default='train')
347 |     parser.add_argument('--dataset_path', default=r'./MVTec')
348 |     parser.add_argument('--category', default='bottle')
349 |     parser.add_argument('--num_epochs', default=1)
350 |     parser.add_argument('--batch_size', default=32)
351 |     parser.add_argument('--load_size', default=256)
352 |     parser.add_argument('--input_size', default=224)
353 |     parser.add_argument('--coreset_sampling_ratio', default=0.001)
354 |     parser.add_argument('--project_root_path', default=r'./test')
355 |     parser.add_argument('--save_src_code', default=True)
356 |     parser.add_argument('--save_anomaly_map', default=True)
357 |     parser.add_argument('--n_neighbors', type=int, default=9)
358 |     args = parser.parse_args()
359 |     return args
360 | 
361 | if __name__ == '__main__':
362 |     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
363 |     args = get_args()
364 |     trainer = pl.Trainer.from_argparse_args(args, default_root_dir=os.path.join(args.project_root_path, args.category), max_epochs=args.num_epochs, gpus=1)
365 |     model = PatchCore(hparams=args)
366 |     if args.phase == 'train':
367 |         trainer.fit(model)
368 |         trainer.test(model)
369 |     elif args.phase == 'test':
370 |         trainer.test(model)
371 | 
372 | 


--------------------------------------------------------------------------------