amazon-science / patchcore-inspection

Apache License 2.0
759 stars 152 forks source link

@authors, can you please elaborate details how to run the repo for any customized dataset. #4

Closed janamoumita1997 closed 2 years ago

janamoumita1997 commented 2 years ago

Actually , I am working on the repo for my current project and will cite this one. Before that please give me detailed procedure for how to run this on customized dataset

Confusezius commented 2 years ago

To include new datasets, you need to write the respective dataloader and put it in src/patchcore/datasets (compare to the mvtec dataloader mvtec.py). Each dataset loader inherits from torch.utils.data.Dataset, and should return the respective image to train (and correspondingly evaluate on) as a dictionary of the form https://github.com/amazon-research/patchcore-inspection/blob/b64be4734cb8295bfbadccf4f6a036b266181e57/src/patchcore/datasets/mvtec.py#L102-L110

Here, the dataloader returns the respective image ("image"), as well as a set of other additional variables useful to evaluate detection and segmentation performance on the test data (such as ground truth masks "mask" and binary variables of whether the input is an anomaly or not "is_anomaly") as well as for visualization the respective image path "image_path".

Once the dataloader is set up, it is referenced in the main training scripts bin/run_patchcore.py, and loaded (with train/(val if needed)/test splits in https://github.com/amazon-research/patchcore-inspection/blob/b64be4734cb8295bfbadccf4f6a036b266181e57/bin/run_patchcore.py#L333-L429

This returns a function handle to the main training part of the script as methods['get_dataloaders'] https://github.com/amazon-research/patchcore-inspection/blob/b64be4734cb8295bfbadccf4f6a036b266181e57/bin/run_patchcore.py#L35 which are then simply used to train the model.

The overall python command on how to call it is shown in the README (with some examples in sample_runs.sh), and for a new dataset would look something like

python bin/run_patchcore.py --gpu 0 --seed 0 --save_patchcore_model --log_group IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0 --log_online --log_project MVTecAD_Results results \
patch_core -b wideresnet50 -le layer2 -le layer3 --faiss_on_gpu --pretrain_embed_dimension 1024  --target_embed_dimension 1024 --anomaly_scorer_num_nn 1 --patchsize 3 sampler -p 0.1 approx_greedy_coreset dataset --resize 256 --imagesize 224 \
>>>Starting from here are the dataset flags and name of dataset (mvtec) to use: "${dataset_flags[@]}" mvtec $datapath

TL;DR: (1) Add pytorch dataloader in src/patchcore/datasets (see mvtec.py for inspiration). (2) Make sure it loads correctly in bin/run_patchcore.py. (3) Run it with commands similar to those in sample_runs.sh and as explained in more detail in the README.

janamoumita1997 commented 2 years ago

Thanks @Confusezius, for your generous reply. But, the issue is the pretrained backbones are not compatible for feature extraction for my dataset. How a model that is trained with my own dataset can be embedded ?

Confusezius commented 2 years ago

For your case, you need to update the backbone you use and where the PatchCore module hooks onto.

The TL;DR in this case is (1) Include your pretrained backbone in src/patchcore/backbones.py with a command that can be used to load it (see other backbones) OR (not really recommended though) directly edit the backbone calls in run_patchcore.py (2) Depending on the architecture, you have to update which layers features are extract from with the -le flags and as done here: https://github.com/amazon-research/patchcore-inspection/blob/b64be4734cb8295bfbadccf4f6a036b266181e57/src/patchcore/patchcore.py#L49-L51.

For a more detailed description, we first have to understand what PatchCore is doing precisely with your pretrained network. For that, note that in the reply above, you can see the parts of the python command that say

patch_core -b wideresnet50 -le layer -le -layer3 ... --pretrain_embed_dimension 1024  --target_embed_dimension 1024

This means that patchcore uses the pretrained wideresnet50 -backbone, and extracts feature maps from the layers (-le = layer_to_extract): [layer2, layer3]. A subsequent step then embeds featuremaps of shapes [h_2 x w_2 x c_2] and [h_3 x w_3 x c_3], respectively, into [h_2 x w_2 x pretrained_embed_dimension] and [h_3 x w_3 x pretrained_embed_dimension]. These are then concatenated by casting the spatially smaller featuremap (in this case layer3) onto the biggest one (in this case layer2), giving [h_2 x w_2 x (pretrained_embed_dimension + pretrained_embed_dimension)]. This is then finally projected down to [h_2 x w_2 x target_emed_dimension].

All this happens here: https://github.com/amazon-research/patchcore-inspection/blob/b64be4734cb8295bfbadccf4f6a036b266181e57/src/patchcore/patchcore.py#L55-L63

You can check out all the details in each Preprocessing and Aggregator class!

With this knowledge, the important thing to understand now is how PatchCore gets the featuremaps it works on. This is done in the layer extraction module here https://github.com/amazon-research/patchcore-inspection/blob/b64be4734cb8295bfbadccf4f6a036b266181e57/src/patchcore/patchcore.py#L49-L51

This extraction module takes in the layer names for a specific backbone (i.e. wideresnet50 and layer2 & layer3), and attaches a forward hook to the final output of each mentioned layer: https://github.com/amazon-research/patchcore-inspection/blob/b64be4734cb8295bfbadccf4f6a036b266181e57/src/patchcore/common.py#L234-L256

Note that layer2 and layer3 are exactly the layernames used in wideresnet50, i.e.

model = timm.create_model('wideresnet50`, pretrained=True)
model.layer2, model.layer3

is well-defined.

So in the case of a novel network, you have to pass in the right layers you want to extract from, sorted by their featuremap size, from largest to smallest (see above; layer2 is larger than layer3). You can also specify more specific locations via e.g. -le layer2.conv2 -le layer3.conv3 for example, assuming that model.layer2.conv2 exists.

Finally, you want to add a loader for your new network in src/patchcore/backbones.py > BACKBONES (see other examples). Alternatively, you can include it directly into bin/run_patchcore.py here: https://github.com/amazon-research/patchcore-inspection/blob/b64be4734cb8295bfbadccf4f6a036b266181e57/bin/run_patchcore.py#L294

Note that this is generally not recommended, as this makes everything a lot less scalable!

To conclude, you can use the python command as

python bin/run_patchcore.py --gpu 0 --seed 0 --save_patchcore_model --log_group IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0 --log_online --log_project MVTecAD_Results results \
>>>Adjust here: patch_core -bnew_network_name -le network_location_1 -le network_location_2 --faiss_on_gpu --pretrain_embed_dimension 1024  --target_embed_dimension 1024 \
--anomaly_scorer_num_nn 1 --patchsize 3 sampler -p 0.1 approx_greedy_coreset dataset --resize 256 --imagesize 224 \
"${dataset_flags[@]}" mvtec $datapath
janamoumita1997 commented 2 years ago

@Confusezius , Thanks for your informative inputs...Got the answer...closing the issue... Thanks @all