├── .gitignore ├── Dockerfile ├── README.md ├── build.bat ├── build.sh ├── configurations.txt ├── detection.py ├── export.bat ├── export.sh ├── model ├── RetinaNetDA.py └── __init__.py ├── model_weights └── RetinaNetDA.pth ├── output └── mitotic-figures.json ├── process.py ├── requirements.txt ├── test.bat ├── test.sh ├── test ├── 007.tiff └── expected_output.json └── util ├── __init__.py ├── data_loader.py ├── nms_WSI.py └── object_detection_helper.py /.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 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | 58 | # Flask stuff: 59 | instance/ 60 | .webassets-cache 61 | 62 | # Scrapy stuff: 63 | .scrapy 64 | 65 | # Sphinx documentation 66 | docs/_build/ 67 | 68 | # PyBuilder 69 | target/ 70 | 71 | # Jupyter Notebook 72 | .ipynb_checkpoints 73 | 74 | # pyenv 75 | .python-version 76 | 77 | # celery beat schedule file 78 | celerybeat-schedule 79 | 80 | # SageMath parsed files 81 | *.sage.py 82 | 83 | # dotenv 84 | .env 85 | 86 | # virtualenv 87 | .venv 88 | venv/ 89 | ENV/ 90 | 91 | # Spyder project settings 92 | .spyderproject 93 | .spyproject 94 | 95 | # Rope project settings 96 | .ropeproject 97 | 98 | # mkdocs documentation 99 | /site 100 | 101 | # mypy 102 | .mypy_cache/ 103 | 104 | # Pycharm 105 | .idea/ 106 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # Edit the base image here, e.g., to use 2 | # TENSORFLOW (https://hub.docker.com/r/tensorflow/tensorflow/) 3 | # or a different PYTORCH (https://hub.docker.com/r/pytorch/pytorch/) base image 4 | FROM pytorch/pytorch:1.6.0-cuda10.1-cudnn7-runtime 5 | 6 | RUN apt-get update 7 | RUN apt-get install -y gcc 8 | 9 | RUN groupadd -r algorithm && useradd -m --no-log-init -r -g algorithm algorithm 10 | 11 | RUN mkdir -p /opt/algorithm /input /output \ 12 | && chown algorithm:algorithm /opt/algorithm /input /output 13 | USER algorithm 14 | 15 | WORKDIR /opt/algorithm 16 | 17 | ENV PATH="/home/algorithm/.local/bin:${PATH}" 18 | 19 | RUN python -m pip install --user -U pip 20 | 21 | # Copy all required files such that they are available within the docker image (code, weights, ...) 22 | COPY --chown=algorithm:algorithm requirements.txt /opt/algorithm/ 23 | 24 | COPY --chown=algorithm:algorithm model/ /opt/algorithm/model/ 25 | COPY --chown=algorithm:algorithm util/ /opt/algorithm/util/ 26 | COPY --chown=algorithm:algorithm model_weights/ /opt/algorithm/checkpoints/ 27 | COPY --chown=algorithm:algorithm process.py /opt/algorithm/ 28 | COPY --chown=algorithm:algorithm detection.py /opt/algorithm/ 29 | 30 | # Install required python packages via pip - you may adapt the requirements.txt to your needs 31 | RUN python -m pip install --user -rrequirements.txt 32 | 33 | # Entrypoint to your python code - executes process.py as a script 34 | ENTRYPOINT python -m process $0 $@ 35 | 36 | ## ALGORITHM LABELS ## 37 | 38 | # These labels are required 39 | LABEL nl.diagnijmegen.rse.algorithm.name=MitosisDetection 40 | 41 | # These labels are required and describe what kind of hardware your algorithm requires to run. 42 | LABEL nl.diagnijmegen.rse.algorithm.hardware.cpu.count=2 43 | LABEL nl.diagnijmegen.rse.algorithm.hardware.cpu.capabilities=() 44 | LABEL nl.diagnijmegen.rse.algorithm.hardware.memory=16G 45 | LABEL nl.diagnijmegen.rse.algorithm.hardware.gpu.count=1 46 | LABEL nl.diagnijmegen.rse.algorithm.hardware.gpu.cuda_compute_capability=6.0 47 | LABEL nl.diagnijmegen.rse.algorithm.hardware.gpu.memory=8G 48 | 49 | 50 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Docker image of reference algorithm for MIDOG 2022 challenge. 2 | 3 | Credits: F. Wilm, K. Breininger, M. Aubreville 4 | 5 | This docker image contains a reference implementation of a domain-adversarial training based on RetinaNet, provided by Frauke Wilm (Friedrich-Alexander-Universität Erlangen-Nürnberg, Germany) for the MIDOG challenge. 6 | 7 | The container shall serve as an example of how we (and the grand-challenge plattform) expect the outputs to look like. At the same time, it serves as a template for you to implement your own algorithm for submission at MIDOG 2022. 8 | 9 | Please note that the MIDOG 2022 docker reference container has changed from the MIDOG 2021 reference container. Main differences are: 10 | - Changed output format (see 2), enabling calculation of mAP as additional metric. 11 | - Updated paths in process.py and test.sh/test.bat to comply with grand-challenge.org's new interface for MIDOG 2022. 12 | 13 | You will have to provide all files to run your model in a docker container. This example may be of help for this. We also provide a quick explanation of how the container works [here](https://www.youtube.com/watch?v=Zkhrwark3bg). 14 | 15 | For reference, you may also want to read the blog post of grand-challenge.org on [how to create an algorithm](https://grand-challenge.org/blogs/create-an-algorithm/). 16 | 17 | ## Content: 18 | 1. [Prerequisites](#prerequisites) 19 | 2. [An overview of the structure of this example](#overview) 20 | 3. [Packing your algorithm into a docker container image](#todocker) 21 | 4. [Building your container](#build) 22 | 5. [Testing your container](#test) 23 | 6. [Generating the bundle for uploading your algorithm](#export) 24 | 25 | ## 1. Prerequisites 26 | 27 | The container is based on docker, so you need to [install docker first](https://www.docker.com/get-started). 28 | 29 | Second, you need to clone this repository: 30 | ``` 31 | git clone https://github.com/DeepPathology/MIDOG_reference_docker 32 | ``` 33 | 34 | You will also need evalutils (provided by grand-challenge): 35 | ``` 36 | pip install evalutils 37 | ``` 38 | 39 | Optional: If you want to have GPU support for local testing, you want to install the [NVIDIA container toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) 40 | 41 | As stated by the grand-challenge team: 42 | >Windows tip: It is highly recommended to install [Windows Subsystem for Linux (WSL)](https://docs.microsoft.com/en-us/windows/wsl/install-win10) to work with Docker on a Linux environment within Windows. Please make sure to install WSL 2 by following the instructions on the same page. In this tutorial, we have used WSL 2 with Ubuntu 18.04 LTS. Also, note that the basic version of WSL 2 does not come with GPU support. Please [watch the official tutorial by Microsoft on installing WSL 2 with GPU support](https://www.youtube.com/watch?v=PdxXlZJiuxA). The alternative is to work purely out of Ubuntu, or any other flavor of Linux. 43 | 44 | ## 2. An overview of the structure of this example 45 | 46 | This example is a RetinaNet implementation, extended by a domain-adversarial branch. 47 | - The main processing (inference) is done in the file [detection.py](detection.py). It provides the class *MyMitosisDetection*, which loads the model and provides the method *process_image()* that takes an individual test image as numpy array as an input and returns the detections on said image. 48 | - The main file that is executed by the container is [process.py](process.py). It imports and instanciates the model (*MyMitosisDetection*). It then loads all images that are part of the test set and processes each of them (using the *process_image()* method). As post-processing, it will also perform a final non-maxima suppression on the image, before creating the return dictionary which contains all individual detected points, which are ultimately stored in the file `/output/mitotic-figures.json`. 49 | 50 | The output file is a dictionary (each input file is processed independently), and has the following format: 51 | 52 | ``` 53 | { 54 | "type": "Multiple points", 55 | "points": [ 56 | { 57 | "point": [ 58 | 0.14647372756903898, 59 | 0.1580733550628604, 60 | 0, 61 | ], 62 | "probability" : 0.534, 63 | "name" : "mitotic figure", 64 | }, 65 | { 66 | "point": [ 67 | 0.11008273935312868, 68 | 0.03707331924495862, 69 | 0, 70 | ] 71 | "probability" : 0.302839283, 72 | "name" : "non-mitotic figure", 73 | } 74 | ], 75 | "version": { 76 | "major": 1, 77 | "minor": 0 78 | } 79 | } 80 | ``` 81 | 82 | Note that each point is described by the following dictionary: 83 | 84 | image 85 | 86 | The field "name" is used to distinguish between above threshold detections and below threshold detections. Please make sure that you find a suitable detection threshold. The below threshold detections are part of the output to calculate the average precision metric. 87 | 88 | **Caution**: This has changed from the MIDOG 2021 docker container and also from earlier versions of this container. If you provide the old format, the evaluation will still work, but will not give you sensible values for the AP metric. 89 | 90 | ## 3. Embedding your algorithm into an algorithm docker container 91 | 92 | We encourage you to adapt this example to your needs and insert your mitosis detection solution. You can adapt the code, remove & code files as needed and adapt parameters, thresholds and other aspects. As discussed above, the main file that is executed by the container is [process.py](process.py). Here, we have marked the most relevant code lines with `TODO`. 93 | 94 | To test this container locally without a docker container, you may set the `execute_in_docker` flag to false - this sets all paths to relative paths. Don't forget to set it back to true when you want to switch back to the docker container setting. 95 | 96 | If you need a different base image to build your container (e.g., Tensorflow instead of Pytorch, or a different version), if you need additional libraries and to make sure that all source files (and weights) are copied to the docker container, you will have to adapt the [Dockerfile](Dockerfile) and the [requirements.txt](requirements.txt) file accordingly. 97 | 98 | Kindly refer to the image below to identify the relevant points: 99 | ![dockerfile_img](https://user-images.githubusercontent.com/43467166/128198999-37dd613d-aeef-41a6-9875-9fdf29db4717.png) 100 | 101 | 102 | ## 4. Building your container 103 | 104 | To test if all dependencies are met, you should run the file `build.bat` (Windows) / `build.sh` (Linux) to build the docker container. Please note that the next step (testing the container) also runs a build, so this step is not mandatory if you are certain that everything is set up correctly. 105 | 106 | ## 5. Testing your container 107 | 108 | To test your container, you should run `test.bat` (on Windows) or `test.sh` (on Linux, might require sudo priviledges). This will run the test image(s) provided in the test folder through your model. It will check them against what you provide in `test/expected_output.json`. Be aware that this will, of course, initially not be equal to the demo detections we put there for testing our reference model. 109 | 110 | ## 6. Generating the bundle for uploading your algorithm 111 | 112 | Finally, you need to run the `export.sh` (Linux) or `export.bat` script to package your docker image. This step creates a file with the extension "tar.gz", which you can then upload to grand-challenge to submit your algorithm. 113 | 114 | ## 7. Creating an "Algorithm" on GrandChallenge and submitting your solution to the MIDOG Challenge 115 | 116 | ** Note: Submission to grand-challenge.org will open on August 5th. ** 117 | 118 | In order to submit your docker container, you first have to add an **Algorithm** entry for your docker container [here] https://midog2022.grand-challenge.org/evaluation/challenge/algorithms/create/. 119 | 120 | Please enter a name for the algorithm: 121 | 122 | image 123 | 124 | After saving, you can add your docker container (you can also overwrite your container here): 125 | 126 | ![uploadcontainer](https://user-images.githubusercontent.com/10051592/128370733-7445e252-a354-4c44-9155-9f232cd9f220.jpg) 127 | 128 | Please note that it can take a while (several minutes) until the container becomes active. You can determine which one is active in the same dialog: 129 | 130 | ![containeractive](https://user-images.githubusercontent.com/10051592/128373241-83102a43-aad7-4457-b068-a6c7cc5a3b98.jpg) 131 | 132 | You can also try out your algorithm. Please note that you will require an image that has the DPI property set in order to use this function. You can use the image test/007.tiff provided as part of this container as test image (it contains mitotic figures). 133 | 134 | ![tryout](https://user-images.githubusercontent.com/10051592/128373614-30b76cf6-2b2d-4d5d-87db-b8c67b47b64f.jpg) 135 | 136 | Finally, you can submit your docker container to MIDOG: 137 | 138 | ![submit_container](https://user-images.githubusercontent.com/10051592/128371715-d8385754-806e-4420-ac5e-4c25cc38112a.jpg) 139 | 140 | ## General remarks 141 | - The training is not done as part of the docker container, so please make sure that you only run inference within the container. 142 | - This image was trained on MIDOG 2021, which had only human breast cancer scanned with various scanners. Do not expect it to have a superb performance on the test sets. 143 | - The official manuscipt has a typo in Table 1. On the XR scanner, the reference approach scored an F1-Score of 0.7826 and thereby outperformed the strong baseline on this scanner. The value was corrected in the manuscript version on arXiv. 144 | 145 | 146 | -------------------------------------------------------------------------------- /build.bat: -------------------------------------------------------------------------------- 1 | docker build -t mitosisdetection . 2 | -------------------------------------------------------------------------------- /build.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )" 3 | 4 | docker build -t mitosisdetection "$SCRIPTPATH" 5 | -------------------------------------------------------------------------------- /configurations.txt: -------------------------------------------------------------------------------- 1 | # network parameters 2 | detect_thresh = 0.64 3 | nms_thresh = 0.4 4 | encoder = create_body(models.resnet18, False, -2) 5 | scales = [0.2, 0.4, 0.6, 0.8, 1.0] 6 | ratios = [1] 7 | sizes = [(64, 64), (32, 32), (16, 16)] 8 | self.model = RetinaNetDA.RetinaNetDA(encoder, n_classes=2, n_domains=4, n_anchors=len(scales) * len(ratios),sizes=[size[0] for size in sizes], chs=128, final_bias=-4., n_conv=3) 9 | 10 | 11 | # normalization parameters 12 | self.model.load_state_dict(torch.load(self.path_model, map_location=self.device)['model']) 13 | self.mean = torch.FloatTensor([0.7481, 0.5692, 0.7225]) # state['data']['normalize']['mean'] 14 | self.std = torch.FloatTensor([0.1759, 0.2284, 0.1792]) # state['data']['normalize']['std'] 15 | -------------------------------------------------------------------------------- /detection.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import torch 3 | from queue import Queue, Empty 4 | from tqdm import tqdm 5 | import numpy as np 6 | import torchvision.transforms as transforms 7 | from util.nms_WSI import nms_patch, nms 8 | from util.object_detection_helper import create_anchors, process_output, rescale_box 9 | from fastai.vision.learner import create_body 10 | from fastai.vision import models 11 | 12 | from model import RetinaNetDA 13 | 14 | 15 | class MyMitosisDetection: 16 | def __init__(self, path_model, size, batchsize, detect_threshold = 0.64, nms_threshold = 0.4): 17 | 18 | # network parameters 19 | self.detect_thresh = detect_threshold 20 | self.nms_thresh = nms_threshold 21 | encoder = create_body(models.resnet18, False, -2) 22 | scales = [0.2, 0.4, 0.6, 0.8, 1.0] 23 | ratios = [1] 24 | sizes = [(64, 64), (32, 32), (16, 16)] 25 | self.model = RetinaNetDA.RetinaNetDA(encoder, n_classes=2, n_domains=4, n_anchors=len(scales) * len(ratios),sizes=[size[0] for size in sizes], chs=128, final_bias=-4., n_conv=3) 26 | self.path_model = path_model 27 | self.size = size 28 | self.batchsize = batchsize 29 | self.mean = None 30 | self.std = None 31 | self.anchors = create_anchors(sizes=sizes, ratios=ratios, scales=scales) 32 | self.device = torch.device('cpu' if not torch.cuda.is_available() else 'cuda') 33 | 34 | def load_model(self): 35 | self.mean = torch.FloatTensor([0.7481, 0.5692, 0.7225]).to(self.device) # state['data']['normalize']['mean'] 36 | self.std = torch.FloatTensor([0.1759, 0.2284, 0.1792]).to(self.device) # state['data']['normalize']['std'] 37 | 38 | if torch.cuda.is_available(): 39 | print("Model loaded on CUDA") 40 | self.model.load_state_dict(torch.load(self.path_model)) 41 | else: 42 | print("Model loaded on CPU") 43 | self.model.load_state_dict(torch.load(self.path_model, map_location='cpu')) 44 | 45 | self.model.to(self.device) 46 | 47 | logging.info("Model loaded. Mean: {} ; Std: {}".format(self.mean, self.std)) 48 | return True 49 | 50 | def process_image(self, input_image): 51 | self.model.eval() 52 | n_patches = 0 53 | queue_patches = Queue() 54 | img_dimensions = input_image.shape 55 | 56 | image_boxes = [] 57 | # create overlapping patches for the whole image 58 | for x in np.arange(0, img_dimensions[1], int(0.9 * self.size)): 59 | for y in np.arange(0, img_dimensions[0], int(0.9 * self.size)): 60 | # last patch shall reach just up to the last pixel 61 | if (x+self.size>img_dimensions[1]): 62 | x = img_dimensions[1]-512 63 | 64 | if (y+self.size>img_dimensions[0]): 65 | y = img_dimensions[0]-512 66 | 67 | queue_patches.put((0, int(x), int(y), input_image)) 68 | n_patches += 1 69 | 70 | 71 | n_batches = int(np.ceil(n_patches / self.batchsize)) 72 | for _ in tqdm(range(n_batches), desc='Processing an image'): 73 | 74 | torch_batch, batch_x, batch_y = self.get_batch(queue_patches) 75 | class_pred_batch, bbox_pred_batch, domain,_ = self.model(torch_batch) 76 | 77 | for b in range(torch_batch.shape[0]): 78 | x_real = batch_x[b] 79 | y_real = batch_y[b] 80 | 81 | cur_class_pred = class_pred_batch[b] 82 | cur_bbox_pred = bbox_pred_batch[b] 83 | cur_patch_boxes = self.postprocess_patch(cur_bbox_pred, cur_class_pred, x_real, y_real) 84 | if len(cur_patch_boxes) > 0: 85 | image_boxes += cur_patch_boxes 86 | 87 | return np.array(image_boxes) 88 | 89 | def get_batch(self, queue_patches): 90 | batch_images = np.zeros((self.batchsize, 3, self.size, self.size)) 91 | batch_x = np.zeros(self.batchsize, dtype=int) 92 | batch_y = np.zeros(self.batchsize, dtype=int) 93 | for i_batch in range(self.batchsize): 94 | if queue_patches.qsize() > 0: 95 | status, batch_x[i_batch], batch_y[i_batch], image = queue_patches.get() 96 | x_start, y_start = int(batch_x[i_batch]), int(batch_y[i_batch]) 97 | 98 | cur_patch = image[y_start:y_start+self.size, x_start:x_start+self.size] / 255. 99 | batch_images[i_batch] = cur_patch.transpose(2, 0, 1)[0:3] 100 | else: 101 | batch_images = batch_images[:i_batch] 102 | batch_x = batch_x[:i_batch] 103 | batch_y = batch_y[:i_batch] 104 | break 105 | torch_batch = torch.from_numpy(batch_images.astype(np.float32, copy=False)).to(self.device) 106 | for p in range(torch_batch.shape[0]): 107 | torch_batch[p] = transforms.Normalize(self.mean, self.std)(torch_batch[p]) 108 | return torch_batch, batch_x, batch_y 109 | 110 | def postprocess_patch(self, cur_bbox_pred, cur_class_pred, x_real, y_real): 111 | cur_patch_boxes = [] 112 | 113 | for clas_pred, bbox_pred in zip(cur_class_pred[None, :, :], cur_bbox_pred[None, :, :], ): 114 | modelOutput = process_output(clas_pred, bbox_pred, self.anchors, self.detect_thresh) 115 | bbox_pred, scores, preds = [modelOutput[x] for x in ['bbox_pred', 'scores', 'preds']] 116 | 117 | if bbox_pred is not None: 118 | # Perform nms per patch to reduce computation effort for the whole image (optional) 119 | to_keep = nms_patch(bbox_pred, scores, self.nms_thresh) 120 | bbox_pred, preds, scores = bbox_pred[to_keep].cpu(), preds[to_keep].cpu(), scores[ 121 | to_keep].cpu() 122 | 123 | t_sz = torch.Tensor([[self.size, self.size]]).float() 124 | 125 | bbox_pred = rescale_box(bbox_pred, t_sz) 126 | 127 | for box, pred, score in zip(bbox_pred, preds, scores): 128 | y_box, x_box = box[:2] 129 | h, w = box[2:4] 130 | 131 | cur_patch_boxes.append( 132 | np.array([x_box + x_real, y_box + y_real, 133 | x_box + x_real + w, y_box + y_real + h, 134 | pred, score])) 135 | 136 | return cur_patch_boxes 137 | 138 | 139 | -------------------------------------------------------------------------------- /export.bat: -------------------------------------------------------------------------------- 1 | call .\build.bat 2 | 3 | docker save mitosisdetection > MitosisDetection.tar 4 | -------------------------------------------------------------------------------- /export.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | ./build.sh 4 | 5 | docker save mitosisdetection | gzip -c > MitosisDetection.tar.gz 6 | -------------------------------------------------------------------------------- /model/RetinaNetDA.py: -------------------------------------------------------------------------------- 1 | from fastai import * 2 | from fastai.vision import * 3 | from fastai.callbacks import * 4 | from fastai.vision.models.unet import _get_sfs_idxs 5 | 6 | # Gradient Reverse Layer 7 | class GradReverse(torch.autograd.Function): 8 | @staticmethod 9 | def forward(ctx, x, alpha): 10 | ctx.alpha = alpha 11 | return x.view_as(x) 12 | 13 | @staticmethod 14 | def backward(ctx, grad_output): 15 | output = grad_output.neg() * ctx.alpha 16 | return output, None 17 | 18 | # export 19 | class LateralUpsampleMerge(nn.Module): 20 | 21 | def __init__(self, ch, ch_lat, hook): 22 | super().__init__() 23 | self.hook = hook 24 | self.conv_lat = conv2d(ch_lat, ch, ks=1, bias=True) 25 | 26 | def forward(self, x): 27 | return self.conv_lat(self.hook.stored) + F.interpolate(x, scale_factor=2) 28 | 29 | class Discriminator(nn.Module): 30 | def __init__(self, size, n_domains, alpha=1.0): 31 | super(Discriminator, self).__init__() 32 | self.alpha = alpha 33 | self.reducer = nn.Sequential( 34 | nn.Conv2d(size, size, kernel_size = (3, 3), bias=False), 35 | nn.BatchNorm2d(size), 36 | nn.ReLU(inplace = True), 37 | nn.Dropout(), 38 | nn.Conv2d(size, size//2, kernel_size = (3, 3), bias=False), 39 | nn.BatchNorm2d(size//2), 40 | nn.ReLU(inplace = True), 41 | nn.Dropout(), 42 | nn.Conv2d(size//2, size//4, kernel_size = (3, 3), bias=False), 43 | nn.BatchNorm2d(size//4), 44 | nn.ReLU(inplace = True), 45 | nn.Dropout(), 46 | nn.AdaptiveAvgPool2d((1, 1)), 47 | )#.cuda() 48 | self.reducer2 = nn.Linear(size//4, n_domains, bias = False)#.cuda() 49 | 50 | for m in self.modules(): 51 | if isinstance(m, nn.Conv2d): 52 | torch.nn.init.normal_(m.weight, mean=0.0, std=0.01) 53 | 54 | def forward(self, x): 55 | x = GradReverse.apply(x, self.alpha) 56 | x = self.reducer(x) 57 | x = torch.flatten(x,1) 58 | x = self.reducer2(x) 59 | return x 60 | 61 | class RetinaNetDA(nn.Module): 62 | "Implements RetinaNet from https://arxiv.org/abs/1708.02002" 63 | 64 | def __init__(self, encoder: nn.Module, n_classes, n_domains, final_bias:float=0., n_conv:float=4, 65 | chs=256, n_anchors=9, flatten=True, sizes=None): 66 | super().__init__() 67 | self.n_classes, self.flatten = n_classes, flatten 68 | imsize = (512, 512) 69 | self.sizes = sizes 70 | sfs_szs, x, hooks = self._model_sizes(encoder, size=imsize) 71 | sfs_idxs = _get_sfs_idxs(sfs_szs) 72 | self.encoder = encoder 73 | self.outputs = hook_outputs(self.encoder[-2:-4:-1]) 74 | self.c5top5 = conv2d(sfs_szs[-1][1], chs, ks=1, bias=True) 75 | self.c5top6 = conv2d(sfs_szs[-1][1], chs, stride=2, bias=True) 76 | self.p6top7 = nn.Sequential(nn.ReLU(), conv2d(chs, chs, stride=2, bias=True)) 77 | self.merges = nn.ModuleList([LateralUpsampleMerge(chs, szs[1], hook) 78 | for szs, hook in zip(sfs_szs[-2:-4:-1], hooks[-2:-4:-1])]) 79 | self.smoothers = nn.ModuleList([conv2d(chs, chs, 3, bias=True) for _ in range(3)]) 80 | self.classifier = self._head_subnet(n_classes, n_anchors, final_bias, chs=chs, n_conv=n_conv) 81 | self.box_regressor = self._head_subnet(4, n_anchors, 0., chs=chs, n_conv=n_conv) 82 | self.n_domains = n_domains 83 | self.d3 = Discriminator(sfs_szs[-3][1], n_domains) 84 | self.d4 = Discriminator(sfs_szs[-2][1], n_domains) 85 | self.d5 = Discriminator(sfs_szs[-1][1], n_domains) 86 | 87 | def _head_subnet(self, n_classes, n_anchors, final_bias=0., n_conv=4, chs=256): 88 | layers = [self._conv2d_relu(chs, chs, bias=True) for _ in range(n_conv)] 89 | layers += [conv2d(chs, n_classes * n_anchors, bias=True)] 90 | layers[-1].bias.data.zero_().add_(final_bias) 91 | layers[-1].weight.data.fill_(0) 92 | return nn.Sequential(*layers) 93 | 94 | def _apply_transpose(self, func, p_states, n_classes): 95 | if not self.flatten: 96 | sizes = [[p.size(0), p.size(2), p.size(3)] for p in p_states] 97 | return [func(p).permute(0, 2, 3, 1).view(*sz, -1, n_classes) for p, sz in zip(p_states, sizes)] 98 | else: 99 | return torch.cat( 100 | [func(p).permute(0, 2, 3, 1).contiguous().view(p.size(0), -1, n_classes) for p in p_states], 1) 101 | 102 | def _model_sizes(self, m: nn.Module, size:tuple=(256,256), full:bool=True) -> Tuple[Sizes,Tensor,Hooks]: 103 | "Passes a dummy input through the model to get the various sizes" 104 | hooks = hook_outputs(m) 105 | ch_in = in_channels(m) 106 | x = torch.zeros(1,ch_in,*size) 107 | x = m.eval()(x) 108 | res = [o.stored.shape for o in hooks] 109 | if not full: hooks.remove() 110 | return res,x,hooks if full else res 111 | 112 | def _conv2d_relu(self, ni:int, nf:int, ks:int=3, stride:int=1, 113 | padding:int=None, bn:bool=False, bias=True) -> nn.Sequential: 114 | "Create a `conv2d` layer with `nn.ReLU` activation and optional(`bn`) `nn.BatchNorm2d`" 115 | layers = [conv2d(ni, nf, ks=ks, stride=stride, padding=padding, bias=bias), nn.ReLU()] 116 | if bn: layers.append(nn.BatchNorm2d(nf)) 117 | return nn.Sequential(*layers) 118 | 119 | def forward(self, x): 120 | c5 = self.encoder(x) 121 | p_states = [self.c5top5(c5.clone()), self.c5top6(c5)] 122 | p_states.append(self.p6top7(p_states[-1])) 123 | for merge in self.merges: 124 | p_states = [merge(p_states[0])] + p_states 125 | for i, smooth in enumerate(self.smoothers[:3]): 126 | p_states[i] = smooth(p_states[i]) 127 | if self.sizes is not None: 128 | p_states = [p_state for p_state in p_states if p_state.size()[-1] in self.sizes] 129 | #d3 = self.d3(self.outputs.stored[1]) 130 | #d4 = self.d4(self.outputs.stored[0]) 131 | d5 = self.d5(c5) 132 | 133 | return [self._apply_transpose(self.classifier, p_states, self.n_classes), 134 | self._apply_transpose(self.box_regressor, p_states, 4), 135 | #d3, 136 | #d4, 137 | d5, 138 | [[p.size(2), p.size(3)] for p in p_states]] 139 | -------------------------------------------------------------------------------- /model/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepMicroscopy/MIDOG_reference_docker/6bdd60294828a1166a819951ba7aa4cb18da7352/model/__init__.py -------------------------------------------------------------------------------- /model_weights/RetinaNetDA.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepMicroscopy/MIDOG_reference_docker/6bdd60294828a1166a819951ba7aa4cb18da7352/model_weights/RetinaNetDA.pth -------------------------------------------------------------------------------- /output/mitotic-figures.json: -------------------------------------------------------------------------------- 1 | {"type": "Multiple points", "points": [{"point": [0.14647372756903898, 0.1580733550628604, 0]}, {"point": [0.11008273935312868, 0.03707331924495862, 0]}], "version": {"major": 1, "minor": 0}} -------------------------------------------------------------------------------- /process.py: -------------------------------------------------------------------------------- 1 | import SimpleITK 2 | from pathlib import Path 3 | 4 | from pandas import DataFrame 5 | import torch 6 | import torchvision 7 | from util.nms_WSI import nms 8 | 9 | from evalutils import DetectionAlgorithm 10 | from evalutils.validators import ( 11 | UniquePathIndicesValidator, 12 | UniqueImagesValidator, 13 | ) 14 | import evalutils 15 | 16 | import json 17 | 18 | from detection import MyMitosisDetection 19 | # TODO: Adapt to MIDOG 2022 reference algos 20 | 21 | # TODO: We have this parameter to adapt the paths between local execution and execution in docker. You can use this flag to switch between these two modes. 22 | execute_in_docker = True 23 | 24 | class Mitosisdetection(DetectionAlgorithm): 25 | def __init__(self): 26 | super().__init__( 27 | validators=dict( 28 | input_image=( 29 | UniqueImagesValidator(), 30 | UniquePathIndicesValidator(), 31 | ) 32 | ), 33 | input_path = Path("/input/images/histopathology-roi-cropout/") if execute_in_docker else Path("./test/"), 34 | output_file = Path("/output/mitotic-figures.json") if execute_in_docker else Path("./output/mitotic-figures.json") 35 | ) 36 | # TODO: This path should lead to your model weights 37 | if execute_in_docker: 38 | path_model = "/opt/algorithm/checkpoints/RetinaNetDA.pth" 39 | else: 40 | path_model = "./model_weights/RetinaNetDA.pth" 41 | 42 | self.size = 512 43 | self.batchsize = 10 44 | self.detect_thresh = 0.64 45 | self.nms_thresh = 0.4 46 | self.level = 0 47 | # TODO: You may adapt this to your model/algorithm here. 48 | ##################################################################################### 49 | # Note: As of MIDOG 2022, the format has changed to enable calculation of the mAP. ## 50 | ##################################################################################### 51 | 52 | # Use NMS threshold as detection threshold for now so we can forward sub-threshold detections to the calculations of the mAP 53 | 54 | self.md = MyMitosisDetection(path_model, self.size, self.batchsize, detect_threshold=self.nms_thresh, nms_threshold=self.nms_thresh) 55 | load_success = self.md.load_model() 56 | if load_success: 57 | print("Successfully loaded model.") 58 | 59 | def save(self): 60 | with open(str(self._output_file), "w") as f: 61 | json.dump(self._case_results[0], f) 62 | 63 | def process_case(self, *, idx, case): 64 | # Load and test the image for this case 65 | input_image, input_image_file_path = self._load_input_image(case=case) 66 | 67 | # Detect and score candidates 68 | scored_candidates = self.predict(input_image=input_image) 69 | 70 | # Write resulting candidates to result.json for this case 71 | return dict(type="Multiple points", points=scored_candidates, version={ "major": 1, "minor": 0 }) 72 | 73 | def predict(self, *, input_image: SimpleITK.Image) -> DataFrame: 74 | # Extract a numpy array with image data from the SimpleITK Image 75 | image_data = SimpleITK.GetArrayFromImage(input_image) 76 | 77 | # TODO: This is the part that you want to adapt to your submission. 78 | with torch.no_grad(): 79 | result_boxes = self.md.process_image(image_data) 80 | 81 | # perform nms per image: 82 | print("All computations done, nms as a last step") 83 | result_boxes = nms(result_boxes, self.nms_thresh) 84 | 85 | candidates = list() 86 | 87 | classnames = ['non-mitotic figure', 'mitotic figure'] 88 | 89 | for i, detection in enumerate(result_boxes): 90 | # our prediction returns x_1, y_1, x_2, y_2, prediction, score -> transform to center coordinates 91 | x_1, y_1, x_2, y_2, prediction, score = detection 92 | coord = tuple(((x_1 + x_2) / 2, (y_1 + y_2) / 2)) 93 | 94 | # For the test set, we expect the coordinates in millimeters - this transformation ensures that the pixel 95 | # coordinates are transformed to mm - if resolution information is available in the .tiff image. If not, 96 | # pixel coordinates are returned. 97 | world_coords = input_image.TransformContinuousIndexToPhysicalPoint( 98 | [c for c in coord] 99 | ) 100 | 101 | 102 | # Expected syntax from evaluation container is: 103 | # x-coordinate(centroid),y-coordinate(centroid),0, detection, score 104 | # where detection should be 1 if score is above threshold and 0 else 105 | candidates.append([*tuple(world_coords),0,int(score>self.detect_thresh), score]) 106 | 107 | result = [{"point": c[0:3], "probability": c[4], "name": classnames[c[3]] } for c in candidates] 108 | return result 109 | 110 | 111 | if __name__ == "__main__": 112 | # loads the image(s), applies DL detection model & saves the result 113 | Mitosisdetection().process() 114 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | evalutils==0.2.4 2 | scikit-learn>=0.21 3 | scipy>=1.2.1 4 | fastai==1.0.61 5 | tqdm 6 | opencv-python==4.5.1.48 7 | torchvision>=0.10.0 8 | torch>=1.6.0 9 | -------------------------------------------------------------------------------- /test.bat: -------------------------------------------------------------------------------- 1 | call .\build.bat 2 | 3 | docker volume create mitosisdetection-output 4 | 5 | docker run --rm^ 6 | --memory=16g^ 7 | -v %~dp0\test\:/input/images/histopathology-roi-cropout/^ 8 | -v mitosisdetection-output:/output/^ 9 | mitosisdetection 10 | 11 | docker run --rm^ 12 | -v mitosisdetection-output:/output/^ 13 | python:3.7-slim cat /output/mitotic-figures.json | python -m json.tool 14 | 15 | docker run --rm^ 16 | -v mitosisdetection-output:/output/^ 17 | -v %~dp0\test\:/input/images/histopathology-roi-cropout/^ 18 | python:3.7-slim python -c "import json, sys; f1 = json.load(open('/output/mitotic-figures.json')); f2 = json.load(open('/input/images/histopathology-roi-cropout/expected_output.json')); sys.exit(f1 != f2);" 19 | 20 | if %ERRORLEVEL% == 0 ( 21 | echo "Tests successfully passed..." 22 | ) else ( 23 | echo "Expected output was not found..." 24 | ) 25 | 26 | docker volume rm mitosisdetection-output 27 | -------------------------------------------------------------------------------- /test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )" 4 | 5 | ./build.sh 6 | 7 | docker volume create mitosisdetection-output 8 | 9 | docker run --rm --gpus all \ 10 | --memory=16g \ 11 | -v $SCRIPTPATH/test/:/input/images/histopathology-roi-cropout/ \ 12 | -v mitosisdetection-output:/output/ \ 13 | mitosisdetection 14 | 15 | docker run --rm \ 16 | -v mitosisdetection-output:/output/ \ 17 | python:3.7-slim cat /output/mitotic-figures.json | python -m json.tool 18 | 19 | docker run --rm \ 20 | -v mitosisdetection-output:/output/ \ 21 | -v $SCRIPTPATH/test/:/input/images/histopathology-roi-cropout/ \ 22 | python:3.7-slim python -c "import json, sys; f1 = json.load(open('/output/mitotic-figures.json')); f2 = json.load(open('/input/images/histopathology-roi-cropout/expected_output.json')); sys.exit(f1 != f2);" 23 | 24 | if [ $? -eq 0 ]; then 25 | echo "Tests successfully passed..." 26 | else 27 | echo "Expected output was not found..." 28 | fi 29 | 30 | docker volume rm mitosisdetection-output 31 | -------------------------------------------------------------------------------- /test/007.tiff: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepMicroscopy/MIDOG_reference_docker/6bdd60294828a1166a819951ba7aa4cb18da7352/test/007.tiff -------------------------------------------------------------------------------- /test/expected_output.json: -------------------------------------------------------------------------------- 1 | { 2 | "type": "Multiple points", 3 | "points": [ 4 | { 5 | "point": [ 6 | 0.14647372756903898, 7 | 0.1580733550628604, 8 | 0 9 | ], 10 | "probability": 0.759753406047821, 11 | "name": "mitotic figure" 12 | }, 13 | { 14 | "point": [ 15 | 0.11008273935312868, 16 | 0.03707331924495862, 17 | 0 18 | ], 19 | "probability": 0.7463231682777405, 20 | "name": "mitotic figure" 21 | }, 22 | { 23 | "point": [ 24 | 0.03536749167233783, 25 | 0.15727730219563738, 26 | 0 27 | ], 28 | "probability": 0.5490298867225647, 29 | "name": "non-mitotic figure" 30 | }, 31 | { 32 | "point": [ 33 | 0.1817274974032021, 34 | 0.20447186503814604, 35 | 0 36 | ], 37 | "probability": 0.4250665307044983, 38 | "name": "non-mitotic figure" 39 | } 40 | ], 41 | "version": { 42 | "major": 1, 43 | "minor": 0 44 | } 45 | } 46 | 47 | -------------------------------------------------------------------------------- /util/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepMicroscopy/MIDOG_reference_docker/6bdd60294828a1166a819951ba7aa4cb18da7352/util/__init__.py -------------------------------------------------------------------------------- /util/data_loader.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | import openslide 3 | import numpy as np 4 | 5 | 6 | class SlideContainer: 7 | 8 | def __init__(self, file: Path, annotations: dict, y, level: int = 0, width: int = 256, height: int = 256, 9 | sample_func: callable = None): 10 | self.file = file 11 | self.slide = openslide.open_slide(str(file)) 12 | self.width = width 13 | self.height = height 14 | self.down_factor = self.slide.level_downsamples[level] 15 | self.y = y 16 | self.annotations = annotations 17 | self.sample_func = sample_func 18 | self.classes = list(set(self.y[1])) 19 | 20 | if level is None: 21 | level = self.slide.level_count - 1 22 | self.level = level 23 | 24 | def get_patch(self, x: int = 0, y: int = 0): 25 | return np.array(self.slide.read_region(location=(int(x * self.down_factor),int(y * self.down_factor)), 26 | level=self.level, size=(self.width, self.height)))[:, :, :3] 27 | 28 | @property 29 | def shape(self): 30 | return self.width, self.height 31 | 32 | def __str__(self): 33 | return 'SlideContainer with:\n sample func: '+str(self.sample_func)+'\n slide:'+str(self.file) 34 | -------------------------------------------------------------------------------- /util/nms_WSI.py: -------------------------------------------------------------------------------- 1 | """ 2 | Non maxima suppression on WSI results 3 | 4 | Uses a kdTree to improve speed. This will only work reasonably for same-sized objects. 5 | 6 | Marc Aubreville, Pattern Recognition Lab, FAU Erlangen-Nürnberg, 2019 7 | """ 8 | 9 | import numpy as np 10 | import torch 11 | from sklearn.neighbors import KDTree 12 | 13 | from util.object_detection_helper import IoU_values 14 | 15 | 16 | def non_max_suppression_by_distance(boxes, scores, radius: float = 25, det_thres=None): 17 | if det_thres is not None: # perform thresholding 18 | to_keep = scores > det_thres 19 | boxes = boxes[to_keep] 20 | scores = scores[to_keep] 21 | 22 | if boxes.shape[-1] == 6: # BBOXES 23 | center_x = boxes[:, 0] + (boxes[:, 2] - boxes[:, 0]) / 2 24 | center_y = boxes[:, 1] + (boxes[:, 3] - boxes[:, 1]) / 2 25 | else: 26 | center_x = boxes[:, 0] 27 | center_y = boxes[:, 1] 28 | 29 | X = np.dstack((center_x, center_y))[0] 30 | tree = KDTree(X) 31 | 32 | sorted_ids = np.argsort(scores)[::-1] 33 | 34 | ids_to_keep = [] 35 | ind = tree.query_radius(X, r=radius) 36 | 37 | while len(sorted_ids) > 0: 38 | ids = sorted_ids[0] 39 | ids_to_keep.append(ids) 40 | sorted_ids = np.delete(sorted_ids, np.in1d(sorted_ids, ind[ids]).nonzero()[0]) 41 | 42 | return boxes[ids_to_keep] 43 | 44 | 45 | def nms(result_boxes, det_thres=None): 46 | arr = np.array(result_boxes) 47 | if arr is not None and isinstance(arr, np.ndarray) and (arr.shape[0] == 0): 48 | return result_boxes 49 | if det_thres is not None: 50 | before = np.sum(arr[:, -1] > det_thres) 51 | if arr.shape[0] > 0: 52 | try: 53 | arr = non_max_suppression_by_distance(arr, arr[:, -1], 25, det_thres) 54 | except: 55 | pass 56 | 57 | result_boxes = arr 58 | 59 | return result_boxes 60 | 61 | 62 | def nms_patch(boxes, scores, thresh=0.5): 63 | idx_sort = scores.argsort(descending=True) 64 | boxes, scores = boxes[idx_sort], scores[idx_sort] 65 | to_keep, indexes = [], torch.LongTensor(np.arange(len(scores))) 66 | 67 | while len(scores) > 0: 68 | to_keep.append(idx_sort[indexes[0]]) 69 | iou_vals = IoU_values(boxes, boxes[:1]).squeeze() 70 | mask_keep = iou_vals <= thresh 71 | if len(mask_keep.nonzero()) == 0: break 72 | boxes, scores, indexes = boxes[mask_keep], scores[mask_keep], indexes[mask_keep] 73 | return torch.LongTensor(to_keep) 74 | -------------------------------------------------------------------------------- /util/object_detection_helper.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import numpy as np 4 | 5 | 6 | def bbox_to_activ(bboxes, anchors, flatten=True): 7 | """Return the target of the model on `anchors` for the `bboxes`.""" 8 | if flatten: 9 | # x and y offsets are normalized by radius 10 | t_centers = (bboxes[..., :2] - anchors[..., :2]) / anchors[..., 2:] 11 | # Radius is given in log scale, relative to anchor radius 12 | t_sizes = torch.log(bboxes[..., 2:] / anchors[..., 2:] + 1e-8) 13 | # Finally, everything is divided by 0.1 (radii by 0.2) 14 | if bboxes.shape[-1] == 4: 15 | return torch.cat([t_centers, t_sizes], -1).div_(bboxes.new_tensor([[0.1, 0.1, 0.2, 0.2]])) 16 | else: 17 | return torch.cat([t_centers, t_sizes], -1).div_(bboxes.new_tensor([[0.1, 0.1, 0.2]])) 18 | 19 | 20 | def create_grid(size): 21 | """Create a grid of a given `size`.""" 22 | H, W = size if isinstance(size, tuple) else (size, size) 23 | grid = torch.FloatTensor(H, W, 2) 24 | linear_points = torch.linspace(-1+1/W, 1-1/W, W) if W > 1 else torch.as_tensor([0.]) 25 | grid[:, :, 1] = torch.ger(torch.ones(H), linear_points).expand_as(grid[:, :, 0]) 26 | linear_points = torch.linspace(-1+1/H, 1-1/H, H) if H > 1 else torch.as_tensor([0.]) 27 | grid[:, :, 0] = torch.ger(linear_points, torch.ones(W)).expand_as(grid[:, :, 1]) 28 | return grid.view(-1, 2) 29 | 30 | 31 | def create_anchors(sizes, ratios, scales, flatten=True): 32 | """Create anchor of `sizes`, `ratios` and `scales`.""" 33 | aspects = [[[s*math.sqrt(r), s*math.sqrt(1/r)] for s in scales] for r in ratios] 34 | aspects = torch.tensor(aspects).view(-1, 2) 35 | anchors = [] 36 | for h, w in sizes: 37 | # 4 here to have the anchors overlap. 38 | sized_aspects = 4 * (aspects * torch.tensor([2/h, 2/w])).unsqueeze(0) 39 | base_grid = create_grid((h, w)).unsqueeze(1) 40 | n, a = base_grid.size(0), aspects.size(0) 41 | ancs = torch.cat([base_grid.expand(n, a, 2), sized_aspects.expand(n, a, 2)], 2) 42 | anchors.append(ancs.view(h, w, a, 4)) 43 | return torch.cat([anc.view(-1, 4) for anc in anchors], 0) if flatten else anchors 44 | 45 | 46 | def rescale_box(bboxes, size: torch.Tensor): 47 | bboxes[:, :2] = bboxes[:, :2] - bboxes[:, 2:] / 2 48 | bboxes[:, :2] = (bboxes[:, :2] + 1) * size / 2 49 | bboxes[:, 2:] = bboxes[:, 2:] * size / 2 50 | bboxes = bboxes.long() 51 | return bboxes 52 | 53 | 54 | def tlbr2cthw(boxes): 55 | """Convert top/left bottom/right format `boxes` to center/size corners.""" 56 | center = (boxes[:, :2] + boxes[:, 2:])/2 57 | sizes = boxes[:, 2:] - boxes[:, :2] 58 | return torch.cat([center, sizes], 1) 59 | 60 | 61 | def encode_class(idxs, n_classes): 62 | target = idxs.new_zeros(len(idxs), n_classes).float() 63 | mask = idxs != 0 64 | i1s = torch.LongTensor(list(range(len(idxs)))) 65 | target[i1s[mask], idxs[mask]-1] = 1 66 | return target 67 | 68 | 69 | def cthw2tlbr(boxes): 70 | """Convert center/size format `boxes` to top/left bottom/right corners.""" 71 | top_left = boxes[:, :2] - boxes[:, 2:]/2 72 | bot_right = boxes[:, :2] + boxes[:, 2:]/2 73 | return torch.cat([top_left, bot_right], 1) 74 | 75 | 76 | def intersection(anchors, targets): 77 | """Compute the sizes of the intersections of `anchors` by `targets`.""" 78 | ancs, tgts = cthw2tlbr(anchors), cthw2tlbr(targets) 79 | a, t = ancs.size(0), tgts.size(0) 80 | ancs, tgts = ancs.unsqueeze(1).expand(a, t, 4), tgts.unsqueeze(0).expand(a, t, 4) 81 | top_left_i = torch.max(ancs[..., :2], tgts[..., :2]) 82 | bot_right_i = torch.min(ancs[..., 2:], tgts[..., 2:]) 83 | sizes = torch.clamp(bot_right_i - top_left_i, min=0) 84 | return sizes[..., 0] * sizes[..., 1] 85 | 86 | 87 | def IoU_values(anchors, targets): 88 | """Compute the IoU values of `anchors` by `targets`.""" 89 | if anchors.shape[-1] == 4: 90 | 91 | inter = intersection(anchors, targets) 92 | anc_sz, tgt_sz = anchors[:, 2] * anchors[:, 3], targets[:, 2] * targets[:, 3] 93 | union = anc_sz.unsqueeze(1) + tgt_sz.unsqueeze(0) - inter 94 | 95 | return inter / (union + 1e-8) 96 | 97 | else: # circular anchors 98 | a, t = anchors.size(0), targets.size(0) 99 | ancs = anchors.unsqueeze(1).expand(a, t, 3) 100 | tgts = targets.unsqueeze(0).expand(a, t, 3) 101 | diff = (ancs[:, :, 0:2] - tgts[:, :, 0:2]) 102 | distances = (diff ** 2).sum(dim=2).sqrt() 103 | radius1 = ancs[..., 2] 104 | radius2 = tgts[..., 2] 105 | acosterm1 = (((distances ** 2) + (radius1 ** 2) - (radius2 ** 2)) / (2 * distances * radius1)).clamp(-1, 106 | 1).acos() 107 | acosterm2 = (((distances ** 2) - (radius1 ** 2) + (radius2 ** 2)) / (2 * distances * radius2)).clamp(-1, 108 | 1).acos() 109 | secondterm = ((radius1 + radius2 - distances) * (distances + radius1 - radius2) * ( 110 | distances + radius1 + radius2) * (distances - radius1 + radius2)).clamp(min=0).sqrt() 111 | 112 | intersec = (radius1 ** 2 * acosterm1) + (radius2 ** 2 * acosterm2) - (0.5 * secondterm) 113 | 114 | union = np.pi * ((radius1 ** 2) + (radius2 ** 2)) - intersec 115 | 116 | return intersec / (union + 1e-8) 117 | 118 | 119 | def match_anchors(anchors, targets, match_thr=0.5, bkg_thr=0.4): 120 | """Match `anchors` to targets. -1 is match to background, -2 is ignore.""" 121 | ious = IoU_values(anchors, targets) 122 | matches = anchors.new(anchors.size(0)).zero_().long() - 2 123 | 124 | if ious.shape[1] > 0: 125 | vals, idxs = torch.max(ious, 1) 126 | matches[vals < bkg_thr] = -1 127 | matches[vals > match_thr] = idxs[vals > match_thr] 128 | return matches 129 | 130 | 131 | def process_output(clas_pred, bbox_pred, anchors, detect_thresh=0.25, use_sigmoid=True): 132 | """Transform predictions to bounding boxes and filter results""" 133 | bbox_pred = activ_to_bbox(bbox_pred, anchors.to(clas_pred.device)) 134 | 135 | if use_sigmoid: 136 | clas_pred = torch.sigmoid(clas_pred) 137 | 138 | clas_pred_orig = clas_pred.clone() 139 | detect_mask = clas_pred.max(1)[0] > detect_thresh 140 | if np.array(detect_mask.cpu()).max() == 0: 141 | return {'bbox_pred': None, 'scores': None, 'preds': None, 'clas_pred': clas_pred, 142 | 'clas_pred_orig': clas_pred_orig, 'detect_mask': detect_mask} 143 | 144 | bbox_pred, clas_pred = bbox_pred[detect_mask], clas_pred[detect_mask] 145 | if bbox_pred.shape[-1] == 4: 146 | bbox_pred = tlbr2cthw(torch.clamp(cthw2tlbr(bbox_pred), min=-1, max=1)) 147 | else: 148 | bbox_pred = bbox_pred # torch.clamp(bbox_pred, min=-1, max=1) 149 | 150 | scores, preds = clas_pred.max(1) 151 | return {'bbox_pred': bbox_pred, 'scores': scores, 'preds': preds, 'clas_pred': clas_pred, 152 | 'clas_pred_orig': clas_pred_orig, 'detect_mask': detect_mask} 153 | 154 | 155 | def activ_to_bbox(acts, anchors, flatten=True): 156 | """Extrapolate bounding boxes on anchors from the model activations.""" 157 | 158 | if flatten: 159 | if anchors.shape[-1] == 4: 160 | acts.mul_(acts.new_tensor([[0.1, 0.1, 0.2, 0.2]])) 161 | centers = anchors[..., 2:] * acts[..., :2] + anchors[..., :2] 162 | sizes = anchors[..., 2:] * torch.exp(acts[..., 2:]) 163 | else: 164 | acts.mul_(acts.new_tensor([[0.1, 0.1, 0.2]])) 165 | centers = anchors[..., 2:] * acts[..., :2] + anchors[..., :2] 166 | sizes = anchors[..., 2:] * torch.exp(acts[..., 2:]) 167 | return torch.cat([centers, sizes], -1) 168 | else: 169 | return [activ_to_bbox(act, anc) for act, anc in zip(acts, anchors)] --------------------------------------------------------------------------------