diff --git a/README.md b/README.md index dcb2794..d2c46f7 100644 --- a/README.md +++ b/README.md @@ -1,139 +1,125 @@ [](https://join.slack.com/t/mdtoolkit/shared_invite/enQtNTQ3MjY2MzE0MDg2LWNjY2I2Njc5MTY0NmM0ZWIxNmQwZDRhYzk2MDdhM2QxYjliYTcwYzhkNTAxYmRkMDA0MjcyNDMyYjllNTZhY2M)


Copyright © German Cancer Research Center (DKFZ), Division of Medical Image Computing (MIC). Please make sure that your usage of this code is in compliance with the code license. ## Overview This is a comprehensive framework for object detection featuring: - 2D + 3D implementations of prevalent object detectors: e.g. Mask R-CNN [1], Retina Net [2], Retina U-Net [3]. - Modular and light-weight structure ensuring sharing of all processing steps (incl. backbone architecture) for comparability of models. - training with bounding box and/or pixel-wise annotations. - dynamic patching and tiling of 2D + 3D images (for training and inference). - weighted consolidation of box predictions across patch-overlaps, ensembles, and dimensions [3]. - monitoring + evaluation simultaneously on object and patient level. - 2D + 3D output visualizations. - integration of COCO mean average precision metric [5]. - integration of MIC-DKFZ batch generators for extensive data augmentation [6]. - easy modification to evaluation of instance segmentation and/or semantic segmentation.
[1] He, Kaiming, et al. "Mask R-CNN" ICCV, 2017
[2] Lin, Tsung-Yi, et al. "Focal Loss for Dense Object Detection" TPAMI, 2018.
[3] Jaeger, Paul et al. "Retina U-Net: Embarrassingly Simple Exploitation of Segmentation Supervision for Medical Object Detection" , 2018 [5] https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py
[6] https://github.com/MIC-DKFZ/batchgenerators

## How to cite this code Please cite the original publication [3]. ## Installation -Setup package in a virtual environment: +Setup package in virtual environment ``` -git clone https://github.com/pfjaeger/medicaldetectiontoolkit.git . +git clone https://github.com/MIC-DKFZ/medicaldetectiontoolkit.git. cd medicaldetectiontoolkit -virtualenv -p python3.6 venv -source venv/bin/activate -pip3 install -e . +virtualenv -p python3.7 mdt +source mdt/bin/activate +python setup.py install ``` +This framework uses two custom mixed C++/CUDA extensions: Non-maximum suppression (NMS) and RoIAlign. Both are adapted from the original pytorch extensions (under torchvision.ops.boxes and ops.roialign). +The extensions are automatically compiled from the provided source files under RegRCNN/custom_extensions with above setup.py. +Note: If you'd like to import the raw extensions (not the wrapper modules), be sure to import torch first. -We use two cuda functions: Non-Maximum Suppression (taken from [pytorch-faster-rcnn](https://github.com/ruotianluo/pytorch-faster-rcnn) and added adaption for 3D) and RoiAlign (taken from [RoiAlign](https://github.com/longcw/RoIAlign.pytorch), fixed according to [this bug report](https://hackernoon.com/how-tensorflows-tf-image-resize-stole-60-days-of-my-life-aba5eb093f35), and added adaption for 3D). In this framework, they come pre-compile for TitanX. If you have a different GPU you need to re-compile these functions: - - -| GPU | arch | -| --- | --- | -| TitanX | sm_52 | -| GTX 960M | sm_50 | -| GTX 1070 | sm_61 | -| GTX 1080 (Ti) | sm_61 | - +Alternatively, you may install the framework via pip by replacing the last line above (python setup.py install) by: ``` -cd cuda_functions/nms_xD/src/cuda/ -nvcc -c -o nms_kernel.cu.o nms_kernel.cu -x cu -Xcompiler -fPIC -arch=[arch] -cd ../../ -python build.py -cd ../ - -cd cuda_functions/roi_align_xD/roi_align/src/cuda/ -nvcc -c -o crop_and_resize_kernel.cu.o crop_and_resize_kernel.cu -x cu -Xcompiler -fPIC -arch=[arch] -cd ../../ -python build.py -cd ../../ +pip install . +pip install -e ./custom_extensions/nms +pip install -e ./custom_extensions/roi_align ``` ## Prepare the Data This framework is meant for you to be able to train models on your own data sets. Two example data loaders are provided in medicaldetectiontoolkit/experiments including thorough documentation to ensure a quick start for your own project. The way I load Data is to have a preprocessing script, which after preprocessing saves the Data of whatever data type into numpy arrays (this is just run once). During training / testing, the data loader then loads these numpy arrays dynamically. (Please note the Data Input side is meant to be customized by you according to your own needs and the provided Data loaders are merely examples: LIDC has a powerful Dataloader that handles 2D/3D inputs and is optimized for patch-based training and inference. Toy-Experiments have a lightweight Dataloader, only handling 2D without patching. The latter makes sense if you want to get familiar with the framework.). ## Execute 1. Set I/O paths, model and training specifics in the configs file: medicaldetectiontoolkit/experiments/your_experiment/configs.py 2. Train the model: ``` python exec.py --mode train --exp_source experiments/my_experiment --exp_dir path/to/experiment/directory ``` This copies snapshots of configs and model to the specified exp_dir, where all outputs will be saved. By default, the data is split into 60% training and 20% validation and 20% testing data to perform a 5-fold cross validation (can be changed to hold-out test set in configs) and all folds will be trained iteratively. In order to train a single fold, specify it using the folds arg: ``` python exec.py --folds 0 1 2 .... # specify any combination of folds [0-4] ``` 3. Run inference: ``` python exec.py --mode test --exp_dir path/to/experiment/directory ``` This runs the prediction pipeline and saves all results to exp_dir. ## Models This framework features all models explored in [3] (implemented in 2D + 3D): The proposed Retina U-Net, a simple but effective Architecture fusing state-of-the-art semantic segmentation with object detection,


also implementations of prevalent object detectors, such as Mask R-CNN, Faster R-CNN+ (Faster R-CNN w\ RoIAlign), Retina Net, U-Faster R-CNN+ (the two stage counterpart of Retina U-Net: Faster R-CNN with auxiliary semantic segmentation), DetU-Net (a U-Net like segmentation architecture with heuristics for object detection.)



## Training annotations This framework features training with pixelwise and/or bounding box annotations. To overcome the issue of box coordinates in data augmentation, we feed the annotation masks through data augmentation (create a pseudo mask, if only bounding box annotations provided) and draw the boxes afterwards.


The framework further handles two types of pixel-wise annotations: 1. A label map with individual ROIs identified by increasing label values, accompanied by a vector containing in each position the class target for the lesion with the corresponding label (for this mode set get_rois_from_seg_flag = False when calling ConvertSegToBoundingBoxCoordinates in your Data Loader). 2. A binary label map. There is only one foreground class and single lesions are not identified. All lesions have the same class target (foreground). In this case the Dataloader runs a Connected Component Labelling algorithm to create processable lesion - class target pairs on the fly (for this mode set get_rois_from_seg_flag = True when calling ConvertSegToBoundingBoxCoordinates in your Data Loader). ## Prediction pipeline This framework provides an inference module, which automatically handles patching of inputs, and tiling, ensembling, and weighted consolidation of output predictions:




## Consolidation of predictions (Weighted Box Clustering) Multiple predictions of the same image (from test time augmentations, tested epochs and overlapping patches), result in a high amount of boxes (or cubes), which need to be consolidated. In semantic segmentation, the final output would typically be obtained by averaging every pixel over all predictions. As described in [3], **weighted box clustering** (WBC) does this for box predictions:





## Visualization / Monitoring By default, loss functions and performance metrics are monitored:



Histograms of matched output predictions for training/validation/testing are plotted per foreground class:


Input images + ground truth annotations + output predictions of a sampled validation abtch are plotted after each epoch (here 2D sampled slice with +-3 neighbouring context slices in channels):


Zoomed into the last two lines of the plot:


## License This framework is published under the [Apache License Version 2.0](LICENSE). diff --git a/cuda_functions/nms_2D/__init__.py b/cuda_functions/nms_2D/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/nms_2D/_ext/__init__.py b/cuda_functions/nms_2D/_ext/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/nms_2D/_ext/nms/__init__.py b/cuda_functions/nms_2D/_ext/nms/__init__.py deleted file mode 100644 index d71786f..0000000 --- a/cuda_functions/nms_2D/_ext/nms/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ - -from torch.utils.ffi import _wrap_function -from ._nms import lib as _lib, ffi as _ffi - -__all__ = [] -def _import_symbols(locals): - for symbol in dir(_lib): - fn = getattr(_lib, symbol) - if callable(fn): - locals[symbol] = _wrap_function(fn, _ffi) - else: - locals[symbol] = fn - __all__.append(symbol) - -_import_symbols(locals()) diff --git a/cuda_functions/nms_2D/_ext/nms/_nms.so b/cuda_functions/nms_2D/_ext/nms/_nms.so deleted file mode 100755 index 1856faf..0000000 Binary files a/cuda_functions/nms_2D/_ext/nms/_nms.so and /dev/null differ diff --git a/cuda_functions/nms_2D/build.py b/cuda_functions/nms_2D/build.py deleted file mode 100644 index 4d9a96b..0000000 --- a/cuda_functions/nms_2D/build.py +++ /dev/null @@ -1,34 +0,0 @@ -import os -import torch -from torch.utils.ffi import create_extension - - -sources = ['src/nms.c'] -headers = ['src/nms.h'] -defines = [] -with_cuda = False - -if torch.cuda.is_available(): - print('Including CUDA code.') - sources += ['src/nms_cuda.c'] - headers += ['src/nms_cuda.h'] - defines += [('WITH_CUDA', None)] - with_cuda = True - -this_file = os.path.dirname(os.path.realpath(__file__)) -print(this_file) -extra_objects = ['src/cuda/nms_kernel.cu.o'] -extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] - -ffi = create_extension( - '_ext.nms', - headers=headers, - sources=sources, - define_macros=defines, - relative_to=__file__, - with_cuda=with_cuda, - extra_objects=extra_objects -) - -if __name__ == '__main__': - ffi.build() diff --git a/cuda_functions/nms_2D/pth_nms.py b/cuda_functions/nms_2D/pth_nms.py deleted file mode 100644 index bfdc29a..0000000 --- a/cuda_functions/nms_2D/pth_nms.py +++ /dev/null @@ -1,39 +0,0 @@ -import torch -from ._ext import nms - - -def nms_gpu(dets, thresh): - """ - dets has to be a tensor - """ - - scores = dets[:, 4] - order = scores.sort(0, descending=True)[1] - dets = dets[order].contiguous() - - keep = torch.LongTensor(dets.size(0)) - num_out = torch.LongTensor(1) - nms.gpu_nms(keep, num_out, dets, thresh) - return order[keep[:num_out[0]].cuda()].contiguous() - - - -def nms_cpu(dets, thresh): - - dets = dets.cpu() - x1 = dets[:, 0] - y1 = dets[:, 1] - x2 = dets[:, 2] - y2 = dets[:, 3] - scores = dets[:, 4] - - areas = (x2 - x1 + 1) * (y2 - y1 + 1) - order = scores.sort(0, descending=True)[1] - # order = torch.from_numpy(np.ascontiguousarray(scores.numpy().argsort()[::-1])).long() - - keep = torch.LongTensor(dets.size(0)) - num_out = torch.LongTensor(1) - nms.cpu_nms(keep, num_out, dets, order, areas, thresh) - - return keep[:num_out[0]] - diff --git a/cuda_functions/nms_2D/src/cuda/nms_kernel.cu b/cuda_functions/nms_2D/src/cuda/nms_kernel.cu deleted file mode 100644 index 1174f22..0000000 --- a/cuda_functions/nms_2D/src/cuda/nms_kernel.cu +++ /dev/null @@ -1,87 +0,0 @@ -// ------------------------------------------------------------------ -// Faster R-CNN -// Copyright (c) 2015 Microsoft -// Licensed under The MIT License [see fast-rcnn/LICENSE for details] -// Written by Shaoqing Ren -// ------------------------------------------------------------------ -#ifdef __cplusplus -extern "C" { -#endif - -#include -#include -#include -#include "nms_kernel.h" - -__device__ inline float devIoU(float const * const a, float const * const b) { - float left = fmaxf(a[0], b[0]), right = fminf(a[2], b[2]); - float top = fmaxf(a[1], b[1]), bottom = fminf(a[3], b[3]); - float width = fmaxf(right - left + 1, 0.f), height = fmaxf(bottom - top + 1, 0.f); - float interS = width * height; - float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1); - float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1); - return interS / (Sa + Sb - interS); -} - -__global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh, - const float *dev_boxes, unsigned long long *dev_mask) { - const int row_start = blockIdx.y; - const int col_start = blockIdx.x; - - // if (row_start > col_start) return; - - const int row_size = - fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock); - const int col_size = - fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock); - - __shared__ float block_boxes[threadsPerBlock * 5]; - if (threadIdx.x < col_size) { - block_boxes[threadIdx.x * 5 + 0] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0]; - block_boxes[threadIdx.x * 5 + 1] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1]; - block_boxes[threadIdx.x * 5 + 2] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2]; - block_boxes[threadIdx.x * 5 + 3] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3]; - block_boxes[threadIdx.x * 5 + 4] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4]; - } - __syncthreads(); - - if (threadIdx.x < row_size) { - const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; - const float *cur_box = dev_boxes + cur_box_idx * 5; - int i = 0; - unsigned long long t = 0; - int start = 0; - if (row_start == col_start) { - start = threadIdx.x + 1; - } - for (i = start; i < col_size; i++) { - if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) { - t |= 1ULL << i; - } - } - const int col_blocks = DIVUP(n_boxes, threadsPerBlock); - dev_mask[cur_box_idx * col_blocks + col_start] = t; - } -} - - -void _nms(int boxes_num, float * boxes_dev, - unsigned long long * mask_dev, float nms_overlap_thresh) { - - dim3 blocks(DIVUP(boxes_num, threadsPerBlock), - DIVUP(boxes_num, threadsPerBlock)); - dim3 threads(threadsPerBlock); - nms_kernel<<>>(boxes_num, - nms_overlap_thresh, - boxes_dev, - mask_dev); -} - -#ifdef __cplusplus -} -#endif diff --git a/cuda_functions/nms_2D/src/cuda/nms_kernel.cu.o b/cuda_functions/nms_2D/src/cuda/nms_kernel.cu.o deleted file mode 100644 index 00135bf..0000000 Binary files a/cuda_functions/nms_2D/src/cuda/nms_kernel.cu.o and /dev/null differ diff --git a/cuda_functions/nms_2D/src/cuda/nms_kernel.h b/cuda_functions/nms_2D/src/cuda/nms_kernel.h deleted file mode 100644 index 2f40582..0000000 --- a/cuda_functions/nms_2D/src/cuda/nms_kernel.h +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef _NMS_KERNEL -#define _NMS_KERNEL - -#ifdef __cplusplus -extern "C" { -#endif - -#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) -int const threadsPerBlock = sizeof(unsigned long long) * 8; - -void _nms(int boxes_num, float * boxes_dev, - unsigned long long * mask_dev, float nms_overlap_thresh); - -#ifdef __cplusplus -} -#endif - -#endif - diff --git a/cuda_functions/nms_2D/src/nms.c b/cuda_functions/nms_2D/src/nms.c deleted file mode 100644 index 4795cc1..0000000 --- a/cuda_functions/nms_2D/src/nms.c +++ /dev/null @@ -1,69 +0,0 @@ -#include -#include - -int cpu_nms(THLongTensor * keep_out, THLongTensor * num_out, THFloatTensor * boxes, THLongTensor * order, THFloatTensor * areas, float nms_overlap_thresh) { - // boxes has to be sorted - THArgCheck(THLongTensor_isContiguous(keep_out), 0, "keep_out must be contiguous"); - THArgCheck(THLongTensor_isContiguous(boxes), 2, "boxes must be contiguous"); - THArgCheck(THLongTensor_isContiguous(order), 3, "order must be contiguous"); - THArgCheck(THLongTensor_isContiguous(areas), 4, "areas must be contiguous"); - // Number of ROIs - long boxes_num = THFloatTensor_size(boxes, 0); - long boxes_dim = THFloatTensor_size(boxes, 1); - - long * keep_out_flat = THLongTensor_data(keep_out); - float * boxes_flat = THFloatTensor_data(boxes); - long * order_flat = THLongTensor_data(order); - float * areas_flat = THFloatTensor_data(areas); - - THByteTensor* suppressed = THByteTensor_newWithSize1d(boxes_num); - THByteTensor_fill(suppressed, 0); - unsigned char * suppressed_flat = THByteTensor_data(suppressed); - - // nominal indices - int i, j; - // sorted indices - int _i, _j; - // temp variables for box i's (the box currently under consideration) - float ix1, iy1, ix2, iy2, iarea; - // variables for computing overlap with box j (lower scoring box) - float xx1, yy1, xx2, yy2; - float w, h; - float inter, ovr; - - long num_to_keep = 0; - for (_i=0; _i < boxes_num; ++_i) { - i = order_flat[_i]; - if (suppressed_flat[i] == 1) { - continue; - } - keep_out_flat[num_to_keep++] = i; - ix1 = boxes_flat[i * boxes_dim]; - iy1 = boxes_flat[i * boxes_dim + 1]; - ix2 = boxes_flat[i * boxes_dim + 2]; - iy2 = boxes_flat[i * boxes_dim + 3]; - iarea = areas_flat[i]; - for (_j = _i + 1; _j < boxes_num; ++_j) { - j = order_flat[_j]; - if (suppressed_flat[j] == 1) { - continue; - } - xx1 = fmaxf(ix1, boxes_flat[j * boxes_dim]); - yy1 = fmaxf(iy1, boxes_flat[j * boxes_dim + 1]); - xx2 = fminf(ix2, boxes_flat[j * boxes_dim + 2]); - yy2 = fminf(iy2, boxes_flat[j * boxes_dim + 3]); - w = fmaxf(0.0, xx2 - xx1 + 1); - h = fmaxf(0.0, yy2 - yy1 + 1); - inter = w * h; - ovr = inter / (iarea + areas_flat[j] - inter); - if (ovr >= nms_overlap_thresh) { - suppressed_flat[j] = 1; - } - } - } - - long *num_out_flat = THLongTensor_data(num_out); - *num_out_flat = num_to_keep; - THByteTensor_free(suppressed); - return 1; -} \ No newline at end of file diff --git a/cuda_functions/nms_2D/src/nms.h b/cuda_functions/nms_2D/src/nms.h deleted file mode 100644 index 25ca0a3..0000000 --- a/cuda_functions/nms_2D/src/nms.h +++ /dev/null @@ -1 +0,0 @@ -int cpu_nms(THLongTensor * keep_out, THLongTensor * num_out, THFloatTensor * boxes, THLongTensor * order, THFloatTensor * areas, float nms_overlap_thresh); \ No newline at end of file diff --git a/cuda_functions/nms_2D/src/nms_cuda.c b/cuda_functions/nms_2D/src/nms_cuda.c deleted file mode 100644 index 5a9a70f..0000000 --- a/cuda_functions/nms_2D/src/nms_cuda.c +++ /dev/null @@ -1,67 +0,0 @@ -// ------------------------------------------------------------------ -// Faster R-CNN -// Copyright (c) 2015 Microsoft -// Licensed under The MIT License [see fast-rcnn/LICENSE for details] -// Written by Shaoqing Ren -// ------------------------------------------------------------------ -#include -#include -#include -#include - -#include "cuda/nms_kernel.h" - - -extern THCState *state; - -int gpu_nms(THLongTensor * keep, THLongTensor* num_out, THCudaTensor * boxes, float nms_overlap_thresh) { - // boxes has to be sorted - THArgCheck(THLongTensor_isContiguous(keep), 0, "boxes must be contiguous"); - THArgCheck(THCudaTensor_isContiguous(state, boxes), 2, "boxes must be contiguous"); - // Number of ROIs - int boxes_num = THCudaTensor_size(state, boxes, 0); - int boxes_dim = THCudaTensor_size(state, boxes, 1); - - float* boxes_flat = THCudaTensor_data(state, boxes); - - const int col_blocks = DIVUP(boxes_num, threadsPerBlock); - THCudaLongTensor * mask = THCudaLongTensor_newWithSize2d(state, boxes_num, col_blocks); - unsigned long long* mask_flat = THCudaLongTensor_data(state, mask); - - _nms(boxes_num, boxes_flat, mask_flat, nms_overlap_thresh); - - THLongTensor * mask_cpu = THLongTensor_newWithSize2d(boxes_num, col_blocks); - THLongTensor_copyCuda(state, mask_cpu, mask); - THCudaLongTensor_free(state, mask); - - unsigned long long * mask_cpu_flat = THLongTensor_data(mask_cpu); - - THLongTensor * remv_cpu = THLongTensor_newWithSize1d(col_blocks); - unsigned long long* remv_cpu_flat = THLongTensor_data(remv_cpu); - THLongTensor_fill(remv_cpu, 0); - - long * keep_flat = THLongTensor_data(keep); - long num_to_keep = 0; - - int i, j; - for (i = 0; i < boxes_num; i++) { - int nblock = i / threadsPerBlock; - int inblock = i % threadsPerBlock; - - if (!(remv_cpu_flat[nblock] & (1ULL << inblock))) { - keep_flat[num_to_keep++] = i; - unsigned long long *p = &mask_cpu_flat[0] + i * col_blocks; - for (j = nblock; j < col_blocks; j++) { - remv_cpu_flat[j] |= p[j]; - } - } - } - - long * num_out_flat = THLongTensor_data(num_out); - * num_out_flat = num_to_keep; - - THLongTensor_free(mask_cpu); - THLongTensor_free(remv_cpu); - - return 1; -} diff --git a/cuda_functions/nms_2D/src/nms_cuda.h b/cuda_functions/nms_2D/src/nms_cuda.h deleted file mode 100644 index 0826111..0000000 --- a/cuda_functions/nms_2D/src/nms_cuda.h +++ /dev/null @@ -1 +0,0 @@ -int gpu_nms(THLongTensor * keep_out, THLongTensor* num_out, THCudaTensor * boxes, float nms_overlap_thresh); \ No newline at end of file diff --git a/cuda_functions/nms_3D/__init__.py b/cuda_functions/nms_3D/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/nms_3D/_ext/__init__.py b/cuda_functions/nms_3D/_ext/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/nms_3D/_ext/nms/__init__.py b/cuda_functions/nms_3D/_ext/nms/__init__.py deleted file mode 100644 index d71786f..0000000 --- a/cuda_functions/nms_3D/_ext/nms/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ - -from torch.utils.ffi import _wrap_function -from ._nms import lib as _lib, ffi as _ffi - -__all__ = [] -def _import_symbols(locals): - for symbol in dir(_lib): - fn = getattr(_lib, symbol) - if callable(fn): - locals[symbol] = _wrap_function(fn, _ffi) - else: - locals[symbol] = fn - __all__.append(symbol) - -_import_symbols(locals()) diff --git a/cuda_functions/nms_3D/_ext/nms/_nms.so b/cuda_functions/nms_3D/_ext/nms/_nms.so deleted file mode 100755 index c8498a0..0000000 Binary files a/cuda_functions/nms_3D/_ext/nms/_nms.so and /dev/null differ diff --git a/cuda_functions/nms_3D/build.py b/cuda_functions/nms_3D/build.py deleted file mode 100644 index 4d9a96b..0000000 --- a/cuda_functions/nms_3D/build.py +++ /dev/null @@ -1,34 +0,0 @@ -import os -import torch -from torch.utils.ffi import create_extension - - -sources = ['src/nms.c'] -headers = ['src/nms.h'] -defines = [] -with_cuda = False - -if torch.cuda.is_available(): - print('Including CUDA code.') - sources += ['src/nms_cuda.c'] - headers += ['src/nms_cuda.h'] - defines += [('WITH_CUDA', None)] - with_cuda = True - -this_file = os.path.dirname(os.path.realpath(__file__)) -print(this_file) -extra_objects = ['src/cuda/nms_kernel.cu.o'] -extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] - -ffi = create_extension( - '_ext.nms', - headers=headers, - sources=sources, - define_macros=defines, - relative_to=__file__, - with_cuda=with_cuda, - extra_objects=extra_objects -) - -if __name__ == '__main__': - ffi.build() diff --git a/cuda_functions/nms_3D/pth_nms.py b/cuda_functions/nms_3D/pth_nms.py deleted file mode 100644 index 3639b5b..0000000 --- a/cuda_functions/nms_3D/pth_nms.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch -from ._ext import nms - - -def nms_gpu(dets, thresh): - """ - dets has to be a tensor - """ - - scores = dets[:, -1] - order = scores.sort(0, descending=True)[1] - dets = dets[order].contiguous() - - keep = torch.LongTensor(dets.size(0)) - num_out = torch.LongTensor(1) - nms.gpu_nms(keep, num_out, dets, thresh) - return order[keep[:num_out[0]].cuda()].contiguous() - - -def nms_cpu(dets, thresh): - - dets = dets.cpu() - x1 = dets[:, 0] - y1 = dets[:, 1] - x2 = dets[:, 2] - y2 = dets[:, 3] - z1 = dets[:, 4] - z2 = dets[:, 5] - scores = dets[:, 6] - areas = (x2 - x1 +1) * (y2 - y1 +1) * (z2 - z1 +1) - order = scores.sort(0, descending=True)[1] - - keep = torch.LongTensor(dets.size(0)) - num_out = torch.LongTensor(1) - nms.cpu_nms(keep, num_out, dets, order, areas, thresh) - - return keep[:num_out[0]] - diff --git a/cuda_functions/nms_3D/src/cuda/nms_kernel.cu b/cuda_functions/nms_3D/src/cuda/nms_kernel.cu deleted file mode 100644 index 5692de8..0000000 --- a/cuda_functions/nms_3D/src/cuda/nms_kernel.cu +++ /dev/null @@ -1,96 +0,0 @@ -// ------------------------------------------------------------------ -// Faster R-CNN -// Copyright (c) 2015 Microsoft -// Licensed under The MIT License [see fast-rcnn/LICENSE for details] -// Written by Shaoqing Ren -// ------------------------------------------------------------------ -#ifdef __cplusplus -extern "C" { -#endif - -#include -#include -#include -#include "nms_kernel.h" - -__device__ inline float devIoU(float const * const a, float const * const b) { - float left = fmaxf(a[0], b[0]), right = fminf(a[2], b[2]); - float top = fmaxf(a[1], b[1]), bottom = fminf(a[3], b[3]); - float front = fmaxf(a[4], b[4]), back = fminf(a[5], b[5]); - - float width = fmaxf(right - left + 1, 0.f), height = fmaxf(bottom - top + 1, 0.f), depth = fmaxf(back - front + 1, 0.f); - float interS = width * height * depth; - float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1) * (a[5] - a[4] + 1); - float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1) * (b[5] - b[4] + 1); - //printf("IoU 3D %f \n", interS / (Sa + Sb - interS)); - - return interS / (Sa + Sb - interS); -} - -__global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh, - const float *dev_boxes, unsigned long long *dev_mask) { - const int row_start = blockIdx.y; - const int col_start = blockIdx.x; - - // if (row_start > col_start) return; - - const int row_size = - fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock); - const int col_size = - fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock); - - __shared__ float block_boxes[threadsPerBlock * 7]; - if (threadIdx.x < col_size) { - block_boxes[threadIdx.x * 7 + 0] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 0]; - block_boxes[threadIdx.x * 7 + 1] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 1]; - block_boxes[threadIdx.x * 7 + 2] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 2]; - block_boxes[threadIdx.x * 7 + 3] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 3]; - block_boxes[threadIdx.x * 7 + 4] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 4]; - block_boxes[threadIdx.x * 7 + 5] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 5]; - block_boxes[threadIdx.x * 7 + 6] = - dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 7 + 6]; - } - __syncthreads(); - - if (threadIdx.x < row_size) { - const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; - const float *cur_box = dev_boxes + cur_box_idx * 7; - int i = 0; - unsigned long long t = 0; - int start = 0; - if (row_start == col_start) { - start = threadIdx.x + 1; - } - for (i = start; i < col_size; i++) { - if (devIoU(cur_box, block_boxes + i * 7) > nms_overlap_thresh) { - t |= 1ULL << i; - } - } - const int col_blocks = DIVUP(n_boxes, threadsPerBlock); - dev_mask[cur_box_idx * col_blocks + col_start] = t; - } -} - - -void _nms(int boxes_num, float * boxes_dev, - unsigned long long * mask_dev, float nms_overlap_thresh) { - - - dim3 blocks(DIVUP(boxes_num, threadsPerBlock), - DIVUP(boxes_num, threadsPerBlock)); - dim3 threads(threadsPerBlock); - nms_kernel<<>>(boxes_num, - nms_overlap_thresh, - boxes_dev, - mask_dev); -} - -#ifdef __cplusplus -} -#endif diff --git a/cuda_functions/nms_3D/src/cuda/nms_kernel.cu.o b/cuda_functions/nms_3D/src/cuda/nms_kernel.cu.o deleted file mode 100644 index ee3ed41..0000000 Binary files a/cuda_functions/nms_3D/src/cuda/nms_kernel.cu.o and /dev/null differ diff --git a/cuda_functions/nms_3D/src/cuda/nms_kernel.h b/cuda_functions/nms_3D/src/cuda/nms_kernel.h deleted file mode 100644 index 2f40582..0000000 --- a/cuda_functions/nms_3D/src/cuda/nms_kernel.h +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef _NMS_KERNEL -#define _NMS_KERNEL - -#ifdef __cplusplus -extern "C" { -#endif - -#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) -int const threadsPerBlock = sizeof(unsigned long long) * 8; - -void _nms(int boxes_num, float * boxes_dev, - unsigned long long * mask_dev, float nms_overlap_thresh); - -#ifdef __cplusplus -} -#endif - -#endif - diff --git a/cuda_functions/nms_3D/src/nms.c b/cuda_functions/nms_3D/src/nms.c deleted file mode 100644 index dd64336..0000000 --- a/cuda_functions/nms_3D/src/nms.c +++ /dev/null @@ -1,74 +0,0 @@ -#include -#include - - -int cpu_nms(THLongTensor * keep_out, THLongTensor * num_out, THFloatTensor * boxes, THLongTensor * order, THFloatTensor * areas, float nms_overlap_thresh) { - // boxes has to be sorted - THArgCheck(THLongTensor_isContiguous(keep_out), 0, "keep_out must be contiguous"); - THArgCheck(THLongTensor_isContiguous(boxes), 2, "boxes must be contiguous"); - THArgCheck(THLongTensor_isContiguous(order), 3, "order must be contiguous"); - THArgCheck(THLongTensor_isContiguous(areas), 4, "areas must be contiguous"); - // Number of ROIs - long boxes_num = THFloatTensor_size(boxes, 0); - long boxes_dim = THFloatTensor_size(boxes, 1); - - long * keep_out_flat = THLongTensor_data(keep_out); - float * boxes_flat = THFloatTensor_data(boxes); - long * order_flat = THLongTensor_data(order); - float * areas_flat = THFloatTensor_data(areas); - - THByteTensor* suppressed = THByteTensor_newWithSize1d(boxes_num); - THByteTensor_fill(suppressed, 0); - unsigned char * suppressed_flat = THByteTensor_data(suppressed); - // nominal indices - int i, j; - // sorted indices - int _i, _j; - // temp variables for box i's (the box currently under consideration) - float ix1, iy1, ix2, iy2, iz1, iz2, iarea; - // variables for computing overlap with box j (lower scoring box) - float xx1, yy1, xx2, yy2, zz1, zz2; - float w, h, d; - float inter, ovr; - - long num_to_keep = 0; - for (_i=0; _i < boxes_num; ++_i) { - i = order_flat[_i]; // from sorted index to nominal index in boxes list. - if (suppressed_flat[i] == 1) { //maybe flag for later. overlapping boxes are surpressed. - continue; - } - keep_out_flat[num_to_keep++] = i; //num to keep is read and then increased. the box index i is saved in keep_out. - ix1 = boxes_flat[i * boxes_dim]; - iy1 = boxes_flat[i * boxes_dim + 1]; - ix2 = boxes_flat[i * boxes_dim + 2]; - iy2 = boxes_flat[i * boxes_dim + 3]; - iz1 = boxes_flat[i * boxes_dim + 4]; - iz2 = boxes_flat[i * boxes_dim + 5]; - iarea = areas_flat[i]; - for (_j = _i + 1; _j < boxes_num; ++_j) { - j = order_flat[_j]; - if (suppressed_flat[j] == 1) { - continue; - } - xx1 = fmaxf(ix1, boxes_flat[j * boxes_dim]); - yy1 = fmaxf(iy1, boxes_flat[j * boxes_dim + 1]); - xx2 = fminf(ix2, boxes_flat[j * boxes_dim + 2]); - yy2 = fminf(iy2, boxes_flat[j * boxes_dim + 3]); - zz1 = fmaxf(iz1, boxes_flat[j * boxes_dim + 4]); - zz2 = fminf(iz2, boxes_flat[j * boxes_dim + 5]); - w = fmaxf(0.0, xx2 - xx1 + 1); - h = fmaxf(0.0, yy2 - yy1 + 1); - d = fmaxf(0.0, zz2 - zz1 + 1); - inter = w * h * d; - ovr = inter / (iarea + areas_flat[j] - inter); - if (ovr >= nms_overlap_thresh) { - suppressed_flat[j] = 1; // can be surpressed because score j < score i (from order: _j = _i + 1 ...) - } - } - } - - long *num_out_flat = THLongTensor_data(num_out); - *num_out_flat = num_to_keep; - THByteTensor_free(suppressed); - return 1; -} \ No newline at end of file diff --git a/cuda_functions/nms_3D/src/nms.h b/cuda_functions/nms_3D/src/nms.h deleted file mode 100644 index d17d9c9..0000000 --- a/cuda_functions/nms_3D/src/nms.h +++ /dev/null @@ -1 +0,0 @@ -int cpu_nms(THLongTensor * keep_out, THLongTensor * num_out, THFloatTensor * boxes, THLongTensor * order, THFloatTensor * areas, float nms_overlap_thresh); diff --git a/cuda_functions/nms_3D/src/nms_cuda.c b/cuda_functions/nms_3D/src/nms_cuda.c deleted file mode 100644 index 5a9a70f..0000000 --- a/cuda_functions/nms_3D/src/nms_cuda.c +++ /dev/null @@ -1,67 +0,0 @@ -// ------------------------------------------------------------------ -// Faster R-CNN -// Copyright (c) 2015 Microsoft -// Licensed under The MIT License [see fast-rcnn/LICENSE for details] -// Written by Shaoqing Ren -// ------------------------------------------------------------------ -#include -#include -#include -#include - -#include "cuda/nms_kernel.h" - - -extern THCState *state; - -int gpu_nms(THLongTensor * keep, THLongTensor* num_out, THCudaTensor * boxes, float nms_overlap_thresh) { - // boxes has to be sorted - THArgCheck(THLongTensor_isContiguous(keep), 0, "boxes must be contiguous"); - THArgCheck(THCudaTensor_isContiguous(state, boxes), 2, "boxes must be contiguous"); - // Number of ROIs - int boxes_num = THCudaTensor_size(state, boxes, 0); - int boxes_dim = THCudaTensor_size(state, boxes, 1); - - float* boxes_flat = THCudaTensor_data(state, boxes); - - const int col_blocks = DIVUP(boxes_num, threadsPerBlock); - THCudaLongTensor * mask = THCudaLongTensor_newWithSize2d(state, boxes_num, col_blocks); - unsigned long long* mask_flat = THCudaLongTensor_data(state, mask); - - _nms(boxes_num, boxes_flat, mask_flat, nms_overlap_thresh); - - THLongTensor * mask_cpu = THLongTensor_newWithSize2d(boxes_num, col_blocks); - THLongTensor_copyCuda(state, mask_cpu, mask); - THCudaLongTensor_free(state, mask); - - unsigned long long * mask_cpu_flat = THLongTensor_data(mask_cpu); - - THLongTensor * remv_cpu = THLongTensor_newWithSize1d(col_blocks); - unsigned long long* remv_cpu_flat = THLongTensor_data(remv_cpu); - THLongTensor_fill(remv_cpu, 0); - - long * keep_flat = THLongTensor_data(keep); - long num_to_keep = 0; - - int i, j; - for (i = 0; i < boxes_num; i++) { - int nblock = i / threadsPerBlock; - int inblock = i % threadsPerBlock; - - if (!(remv_cpu_flat[nblock] & (1ULL << inblock))) { - keep_flat[num_to_keep++] = i; - unsigned long long *p = &mask_cpu_flat[0] + i * col_blocks; - for (j = nblock; j < col_blocks; j++) { - remv_cpu_flat[j] |= p[j]; - } - } - } - - long * num_out_flat = THLongTensor_data(num_out); - * num_out_flat = num_to_keep; - - THLongTensor_free(mask_cpu); - THLongTensor_free(remv_cpu); - - return 1; -} diff --git a/cuda_functions/nms_3D/src/nms_cuda.h b/cuda_functions/nms_3D/src/nms_cuda.h deleted file mode 100644 index 08bf147..0000000 --- a/cuda_functions/nms_3D/src/nms_cuda.h +++ /dev/null @@ -1 +0,0 @@ -int gpu_nms(THLongTensor * keep_out, THLongTensor* num_out, THCudaTensor * boxes, float nms_overlap_thresh); diff --git a/cuda_functions/roi_align_2D/__init__.py b/cuda_functions/roi_align_2D/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/roi_align_2D/roi_align/__init__.py b/cuda_functions/roi_align_2D/roi_align/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/roi_align_2D/roi_align/_ext/__init__.py b/cuda_functions/roi_align_2D/roi_align/_ext/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/roi_align_2D/roi_align/_ext/crop_and_resize/__init__.py b/cuda_functions/roi_align_2D/roi_align/_ext/crop_and_resize/__init__.py deleted file mode 100644 index 4486c09..0000000 --- a/cuda_functions/roi_align_2D/roi_align/_ext/crop_and_resize/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ - -from torch.utils.ffi import _wrap_function -from ._crop_and_resize import lib as _lib, ffi as _ffi - -__all__ = [] -def _import_symbols(locals): - for symbol in dir(_lib): - fn = getattr(_lib, symbol) - if callable(fn): - locals[symbol] = _wrap_function(fn, _ffi) - else: - locals[symbol] = fn - __all__.append(symbol) - -_import_symbols(locals()) diff --git a/cuda_functions/roi_align_2D/roi_align/_ext/crop_and_resize/_crop_and_resize.so b/cuda_functions/roi_align_2D/roi_align/_ext/crop_and_resize/_crop_and_resize.so deleted file mode 100755 index e852f11..0000000 Binary files a/cuda_functions/roi_align_2D/roi_align/_ext/crop_and_resize/_crop_and_resize.so and /dev/null differ diff --git a/cuda_functions/roi_align_2D/roi_align/build.py b/cuda_functions/roi_align_2D/roi_align/build.py deleted file mode 100755 index 3798d82..0000000 --- a/cuda_functions/roi_align_2D/roi_align/build.py +++ /dev/null @@ -1,40 +0,0 @@ -import os -import torch -from torch.utils.ffi import create_extension - - -sources = ['src/crop_and_resize.c'] -headers = ['src/crop_and_resize.h'] -defines = [] -with_cuda = False - -extra_objects = [] -if torch.cuda.is_available(): - print('Including CUDA code.') - sources += ['src/crop_and_resize_gpu.c'] - headers += ['src/crop_and_resize_gpu.h'] - defines += [('WITH_CUDA', None)] - extra_objects += ['src/cuda/crop_and_resize_kernel.cu.o'] - with_cuda = True - -extra_compile_args = ['-fopenmp', '-std=c99'] - -this_file = os.path.dirname(os.path.realpath(__file__)) -print(this_file) -sources = [os.path.join(this_file, fname) for fname in sources] -headers = [os.path.join(this_file, fname) for fname in headers] -extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] - -ffi = create_extension( - '_ext.crop_and_resize', - headers=headers, - sources=sources, - define_macros=defines, - relative_to=__file__, - with_cuda=with_cuda, - extra_objects=extra_objects, - extra_compile_args=extra_compile_args -) - -if __name__ == '__main__': - ffi.build() diff --git a/cuda_functions/roi_align_2D/roi_align/crop_and_resize.py b/cuda_functions/roi_align_2D/roi_align/crop_and_resize.py deleted file mode 100755 index 4291ae4..0000000 --- a/cuda_functions/roi_align_2D/roi_align/crop_and_resize.py +++ /dev/null @@ -1,66 +0,0 @@ -import math -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.autograd import Function - -from ._ext import crop_and_resize as _backend - - -class CropAndResizeFunction(Function): - - def __init__(self, crop_height, crop_width, extrapolation_value=0): - self.crop_height = crop_height - self.crop_width = crop_width - self.extrapolation_value = extrapolation_value - - def forward(self, image, boxes, box_ind): - crops = torch.zeros_like(image) - if image.is_cuda: - _backend.crop_and_resize_gpu_forward( - image, boxes, box_ind, - self.extrapolation_value, self.crop_height, self.crop_width, crops) - else: - _backend.crop_and_resize_forward( - image, boxes, box_ind, - self.extrapolation_value, self.crop_height, self.crop_width, crops) - - # save for backward - self.im_size = image.size() - self.save_for_backward(boxes, box_ind) - - return crops - - def backward(self, grad_outputs): - boxes, box_ind = self.saved_tensors - - grad_outputs = grad_outputs.contiguous() - grad_image = torch.zeros_like(grad_outputs).resize_(*self.im_size) - - if grad_outputs.is_cuda: - _backend.crop_and_resize_gpu_backward( - grad_outputs, boxes, box_ind, grad_image - ) - else: - _backend.crop_and_resize_backward( - grad_outputs, boxes, box_ind, grad_image - ) - - return grad_image, None, None - - -class CropAndResize(nn.Module): - """ - Crop and resize ported from tensorflow - See more details on https://www.tensorflow.org/api_docs/python/tf/image/crop_and_resize - """ - - def __init__(self, crop_height, crop_width, extrapolation_value=0): - super(CropAndResize, self).__init__() - - self.crop_height = crop_height - self.crop_width = crop_width - self.extrapolation_value = extrapolation_value - - def forward(self, image, boxes, box_ind): - return CropAndResizeFunction(self.crop_height, self.crop_width, self.extrapolation_value)(image, boxes, box_ind) diff --git a/cuda_functions/roi_align_2D/roi_align/roi_align.py b/cuda_functions/roi_align_2D/roi_align/roi_align.py deleted file mode 100644 index 6931539..0000000 --- a/cuda_functions/roi_align_2D/roi_align/roi_align.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -from torch import nn - -from .crop_and_resize import CropAndResizeFunction, CropAndResize - - -class RoIAlign(nn.Module): - - def __init__(self, crop_height, crop_width, extrapolation_value=0, transform_fpcoor=True): - super(RoIAlign, self).__init__() - - self.crop_height = crop_height - self.crop_width = crop_width - self.extrapolation_value = extrapolation_value - self.transform_fpcoor = transform_fpcoor - - def forward(self, featuremap, boxes, box_ind): - """ - RoIAlign based on crop_and_resize. - See more details on https://github.com/ppwwyyxx/tensorpack/blob/6d5ba6a970710eaaa14b89d24aace179eb8ee1af/examples/FasterRCNN/model.py#L301 - :param featuremap: NxCxHxW - :param boxes: Mx4 float box with (x1, y1, x2, y2) **without normalization** - :param box_ind: M - :return: MxCxoHxoW - """ - x1, y1, x2, y2 = torch.split(boxes, 1, dim=1) - image_height, image_width = featuremap.size()[2:4] - - if self.transform_fpcoor: - spacing_w = (x2 - x1) / float(self.crop_width) - spacing_h = (y2 - y1) / float(self.crop_height) - - nx0 = (x1 + spacing_w / 2 - 0.5) / float(image_width - 1) - ny0 = (y1 + spacing_h / 2 - 0.5) / float(image_height - 1) - nw = spacing_w * float(self.crop_width - 1) / float(image_width - 1) - nh = spacing_h * float(self.crop_height - 1) / float(image_height - 1) - - boxes = torch.cat((ny0, nx0, ny0 + nh, nx0 + nw), 1) - else: - x1 = x1 / float(image_width - 1) - x2 = x2 / float(image_width - 1) - y1 = y1 / float(image_height - 1) - y2 = y2 / float(image_height - 1) - boxes = torch.cat((y1, x1, y2, x2), 1) - - boxes = boxes.detach().contiguous() - box_ind = box_ind.detach() - return CropAndResizeFunction(self.crop_height, self.crop_width, self.extrapolation_value)(featuremap, boxes, box_ind) diff --git a/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize.c b/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize.c deleted file mode 100644 index a5ff973..0000000 --- a/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize.c +++ /dev/null @@ -1,269 +0,0 @@ -#include -#include -#include - - -void CropAndResizePerBox( - const float * image_data, - const int batch_size, - const int depth, - const int image_height, - const int image_width, - - const float * boxes_data, - const int * box_index_data, - const int start_box, - const int limit_box, - - float * corps_data, - const int crop_height, - const int crop_width, - const float extrapolation_value -) { - const int image_channel_elements = image_height * image_width; - const int image_elements = depth * image_channel_elements; - - const int channel_elements = crop_height * crop_width; - const int crop_elements = depth * channel_elements; - - int b; - #pragma omp parallel for - for (b = start_box; b < limit_box; ++b) { - const float * box = boxes_data + b * 4; - const float y1 = box[0]; - const float x1 = box[1]; - const float y2 = box[2]; - const float x2 = box[3]; - - const int b_in = box_index_data[b]; - if (b_in < 0 || b_in >= batch_size) { - printf("Error: batch_index %d out of range [0, %d)\n", b_in, batch_size); - exit(-1); - } - - const float height_scale = - (crop_height > 1) - ? (y2 - y1) * (image_height - 1) / (crop_height - 1) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) - : 0; - - for (int y = 0; y < crop_height; ++y) - { - const float in_y = (crop_height > 1) - ? y1 * (image_height - 1) + y * height_scale - : 0.5 * (y1 + y2) * (image_height - 1); - - if (in_y < 0 || in_y > image_height - 1) - { - for (int x = 0; x < crop_width; ++x) - { - for (int d = 0; d < depth; ++d) - { - // crops(b, y, x, d) = extrapolation_value; - corps_data[crop_elements * b + channel_elements * d + y * crop_width + x] = extrapolation_value; - } - } - continue; - } - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - for (int x = 0; x < crop_width; ++x) - { - const float in_x = (crop_width > 1) - ? x1 * (image_width - 1) + x * width_scale - : 0.5 * (x1 + x2) * (image_width - 1); - if (in_x < 0 || in_x > image_width - 1) - { - for (int d = 0; d < depth; ++d) - { - corps_data[crop_elements * b + channel_elements * d + y * crop_width + x] = extrapolation_value; - } - continue; - } - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - for (int d = 0; d < depth; ++d) - { - const float *pimage = image_data + b_in * image_elements + d * image_channel_elements; - - const float top_left = pimage[top_y_index * image_width + left_x_index]; - const float top_right = pimage[top_y_index * image_width + right_x_index]; - const float bottom_left = pimage[bottom_y_index * image_width + left_x_index]; - const float bottom_right = pimage[bottom_y_index * image_width + right_x_index]; - - const float top = top_left + (top_right - top_left) * x_lerp; - const float bottom = - bottom_left + (bottom_right - bottom_left) * x_lerp; - - corps_data[crop_elements * b + channel_elements * d + y * crop_width + x] = top + (bottom - top) * y_lerp; - } - } // end for x - } // end for y - } // end for b - -} - - -void crop_and_resize_forward( - THFloatTensor * image, - THFloatTensor * boxes, // [y1, x1, y2, x2] - THIntTensor * box_index, // range in [0, batch_size) - const float extrapolation_value, - const int crop_height, - const int crop_width, - THFloatTensor * crops -) { - //const int batch_size = image->size[0]; - //const int depth = image->size[1]; - //const int image_height = image->size[2]; - //const int image_width = image->size[3]; - - //const int num_boxes = boxes->size[0]; - - const int batch_size = THFloatTensor_size(image, 0); - const int depth = THFloatTensor_size(image, 1); - const int image_height = THFloatTensor_size(image, 2); - const int image_width = THFloatTensor_size(image, 3); - - const int num_boxes = THFloatTensor_size(boxes, 0); - - // init output space - THFloatTensor_resize4d(crops, num_boxes, depth, crop_height, crop_width); - THFloatTensor_zero(crops); - - // crop_and_resize for each box - CropAndResizePerBox( - THFloatTensor_data(image), - batch_size, - depth, - image_height, - image_width, - - THFloatTensor_data(boxes), - THIntTensor_data(box_index), - 0, - num_boxes, - - THFloatTensor_data(crops), - crop_height, - crop_width, - extrapolation_value - ); - -} - - -void crop_and_resize_backward( - THFloatTensor * grads, - THFloatTensor * boxes, // [y1, x1, y2, x2] - THIntTensor * box_index, // range in [0, batch_size) - THFloatTensor * grads_image // resize to [bsize, c, hc, wc] -) -{ - // shape - //const int batch_size = grads_image->size[0]; - //const int depth = grads_image->size[1]; - //const int image_height = grads_image->size[2]; - //const int image_width = grads_image->size[3]; - - //const int num_boxes = grads->size[0]; - //const int crop_height = grads->size[2]; - //const int crop_width = grads->size[3]; - - const int batch_size = THFloatTensor_size(grads_image, 0); - const int depth = THFloatTensor_size(grads_image, 1); - const int image_height = THFloatTensor_size(grads_image, 2); - const int image_width = THFloatTensor_size(grads_image, 3); - - const int num_boxes = THFloatTensor_size(grads, 0); - const int crop_height = THFloatTensor_size(grads,2); - const int crop_width = THFloatTensor_size(grads,3); - - - // n_elements - const int image_channel_elements = image_height * image_width; - const int image_elements = depth * image_channel_elements; - - const int channel_elements = crop_height * crop_width; - const int crop_elements = depth * channel_elements; - - // init output space - THFloatTensor_zero(grads_image); - - // data pointer - const float * grads_data = THFloatTensor_data(grads); - const float * boxes_data = THFloatTensor_data(boxes); - const int * box_index_data = THIntTensor_data(box_index); - float * grads_image_data = THFloatTensor_data(grads_image); - - for (int b = 0; b < num_boxes; ++b) { - const float * box = boxes_data + b * 4; - const float y1 = box[0]; - const float x1 = box[1]; - const float y2 = box[2]; - const float x2 = box[3]; - - const int b_in = box_index_data[b]; - if (b_in < 0 || b_in >= batch_size) { - printf("Error: batch_index %d out of range [0, %d)\n", b_in, batch_size); - exit(-1); - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height - 1) / (crop_height - 1) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) - : 0; - - for (int y = 0; y < crop_height; ++y) - { - const float in_y = (crop_height > 1) - ? y1 * (image_height - 1) + y * height_scale - : 0.5 * (y1 + y2) * (image_height - 1); - if (in_y < 0 || in_y > image_height - 1) - { - continue; - } - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - for (int x = 0; x < crop_width; ++x) - { - const float in_x = (crop_width > 1) - ? x1 * (image_width - 1) + x * width_scale - : 0.5 * (x1 + x2) * (image_width - 1); - if (in_x < 0 || in_x > image_width - 1) - { - continue; - } - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - for (int d = 0; d < depth; ++d) - { - float *pimage = grads_image_data + b_in * image_elements + d * image_channel_elements; - const float grad_val = grads_data[crop_elements * b + channel_elements * d + y * crop_width + x]; - - const float dtop = (1 - y_lerp) * grad_val; - pimage[top_y_index * image_width + left_x_index] += (1 - x_lerp) * dtop; - pimage[top_y_index * image_width + right_x_index] += x_lerp * dtop; - - const float dbottom = y_lerp * grad_val; - pimage[bottom_y_index * image_width + left_x_index] += (1 - x_lerp) * dbottom; - pimage[bottom_y_index * image_width + right_x_index] += x_lerp * dbottom; - } // end d - } // end x - } // end y - } // end b -} \ No newline at end of file diff --git a/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize.h b/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize.h deleted file mode 100644 index d494865..0000000 --- a/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize.h +++ /dev/null @@ -1,16 +0,0 @@ -void crop_and_resize_forward( - THFloatTensor * image, - THFloatTensor * boxes, // [y1, x1, y2, x2] - THIntTensor * box_index, // range in [0, batch_size) - const float extrapolation_value, - const int crop_height, - const int crop_width, - THFloatTensor * crops -); - -void crop_and_resize_backward( - THFloatTensor * grads, - THFloatTensor * boxes, // [y1, x1, y2, x2] - THIntTensor * box_index, // range in [0, batch_size) - THFloatTensor * grads_image // resize to [bsize, c, hc, wc] -); \ No newline at end of file diff --git a/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize_gpu.c b/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize_gpu.c deleted file mode 100644 index dd347c6..0000000 --- a/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize_gpu.c +++ /dev/null @@ -1,68 +0,0 @@ -#include -#include "cuda/crop_and_resize_kernel.h" - -extern THCState *state; - - -void crop_and_resize_gpu_forward( - THCudaTensor * image, - THCudaTensor * boxes, // [y1, x1, y2, x2] - THCudaIntTensor * box_index, // range in [0, batch_size) - const float extrapolation_value, - const int crop_height, - const int crop_width, - THCudaTensor * crops -) { - const int batch_size = THCudaTensor_size(state, image, 0); - const int depth = THCudaTensor_size(state, image, 1); - const int image_height = THCudaTensor_size(state, image, 2); - const int image_width = THCudaTensor_size(state, image, 3); - - const int num_boxes = THCudaTensor_size(state, boxes, 0); - - // init output space - THCudaTensor_resize4d(state, crops, num_boxes, depth, crop_height, crop_width); - THCudaTensor_zero(state, crops); - cudaStream_t stream = THCState_getCurrentStream(state); - CropAndResizeLaucher( - THCudaTensor_data(state, image), - THCudaTensor_data(state, boxes), - THCudaIntTensor_data(state, box_index), - num_boxes, batch_size, image_height, image_width, - crop_height, crop_width, depth, extrapolation_value, - THCudaTensor_data(state, crops), - stream - ); -} - - -void crop_and_resize_gpu_backward( - THCudaTensor * grads, - THCudaTensor * boxes, // [y1, x1, y2, x2] - THCudaIntTensor * box_index, // range in [0, batch_size) - THCudaTensor * grads_image // resize to [bsize, c, hc, wc] -) { - // shape - const int batch_size = THCudaTensor_size(state, grads_image, 0); - const int depth = THCudaTensor_size(state, grads_image, 1); - const int image_height = THCudaTensor_size(state, grads_image, 2); - const int image_width = THCudaTensor_size(state, grads_image, 3); - - const int num_boxes = THCudaTensor_size(state, grads, 0); - const int crop_height = THCudaTensor_size(state, grads, 2); - const int crop_width = THCudaTensor_size(state, grads, 3); - - // init output space - THCudaTensor_zero(state, grads_image); - - cudaStream_t stream = THCState_getCurrentStream(state); - CropAndResizeBackpropImageLaucher( - THCudaTensor_data(state, grads), - THCudaTensor_data(state, boxes), - THCudaIntTensor_data(state, box_index), - num_boxes, batch_size, image_height, image_width, - crop_height, crop_width, depth, - THCudaTensor_data(state, grads_image), - stream - ); -} \ No newline at end of file diff --git a/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize_gpu.h b/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize_gpu.h deleted file mode 100644 index c2a64cf..0000000 --- a/cuda_functions/roi_align_2D/roi_align/src/crop_and_resize_gpu.h +++ /dev/null @@ -1,16 +0,0 @@ -void crop_and_resize_gpu_forward( - THCudaTensor * image, - THCudaTensor * boxes, // [y1, x1, y2, x2] - THCudaIntTensor * box_index, // range in [0, batch_size) - const float extrapolation_value, - const int crop_height, - const int crop_width, - THCudaTensor * crops -); - -void crop_and_resize_gpu_backward( - THCudaTensor * grads, - THCudaTensor * boxes, // [y1, x1, y2, x2] - THCudaIntTensor * box_index, // range in [0, batch_size) - THCudaTensor * grads_image // resize to [bsize, c, hc, wc] -); \ No newline at end of file diff --git a/cuda_functions/roi_align_2D/roi_align/src/cuda/backup.cu b/cuda_functions/roi_align_2D/roi_align/src/cuda/backup.cu deleted file mode 100644 index 3a1ab8b..0000000 --- a/cuda_functions/roi_align_2D/roi_align/src/cuda/backup.cu +++ /dev/null @@ -1,243 +0,0 @@ -#include -#include -#include "crop_and_resize_kernel.h" - -#define CUDA_1D_KERNEL_LOOP(i, n) \ -for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \ - i += blockDim.x * gridDim.x) - - -__global__ -void CropAndResizeKernel( - const int nthreads, const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float extrapolation_value, float *crops_ptr) -{ - CUDA_1D_KERNEL_LOOP(out_idx, nthreads) - { - // NHWC: out_idx = d + depth * (w + crop_width * (h + crop_height * b)) - // NCHW: out_idx = w + crop_width * (h + crop_height * (d + depth * b)) - int idx = out_idx; - const int x = idx % crop_width; - idx /= crop_width; - const int y = idx % crop_height; - idx /= crop_height; - const int d = idx % depth; - const int b = idx / depth; - - const float y1 = boxes_ptr[b * 4]; - const float x1 = boxes_ptr[b * 4 + 1]; - const float y2 = boxes_ptr[b * 4 + 2]; - const float x2 = boxes_ptr[b * 4 + 3]; - - // printf("INIT CUDA SCRIPT %f \n", idx); - - const int b_in = box_ind_ptr[b]; - if (b_in < 0 || b_in >= batch) - { - continue; - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height - 1) / (crop_height - 1) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) : 0; - - const float in_y = (crop_height > 1) - ? y1 * (image_height - 1) + y * height_scale - : 0.5 * (y1 + y2) * (image_height - 1); - if (in_y < 0 || in_y > image_height - 1) - { - crops_ptr[out_idx] = extrapolation_value; - continue; - } - - const float in_x = (crop_width > 1) - ? x1 * (image_width - 1) + x * width_scale - : 0.5 * (x1 + x2) * (image_width - 1); - if (in_x < 0 || in_x > image_width - 1) - { - crops_ptr[out_idx] = extrapolation_value; - continue; - } - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - const float *pimage = image_ptr + (b_in * depth + d) * image_height * image_width; - const float top_left = pimage[top_y_index * image_width + left_x_index]; - const float top_right = pimage[top_y_index * image_width + right_x_index]; - const float bottom_left = pimage[bottom_y_index * image_width + left_x_index]; - const float bottom_right = pimage[bottom_y_index * image_width + right_x_index]; - // if (top_left == 0){ - // const float top = top_right} - // elif (top_right == 0){ - // const float top = top_left} - // else{ - const float top = top_left + (top_right - top_left) * x_lerp; - //} - - //if (bottom_left == 0){ - // const float bottom = bottom_right} - // elif (bottom_right == 0){ - // const float bottom = bottom_left} - // else{ - const float bottom = bottom_left + (bottom_right - bottom_left) * x_lerp; - //} - - //if (top == 0){ - // crops_ptr[out_idx] = bottom } - // elif (bottom == 0){ - // crops_ptr[out_idx] = top - //} - // else{ - crops_ptr[out_idx] = top + (bottom - top) * y_lerp; - //} - } -} - -__global__ -void CropAndResizeBackpropImageKernel( - const int nthreads, const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float *grads_image_ptr) -{ - CUDA_1D_KERNEL_LOOP(out_idx, nthreads) - { - // NHWC: out_idx = d + depth * (w + crop_width * (h + crop_height * b)) - // NCHW: out_idx = w + crop_width * (h + crop_height * (d + depth * b)) - int idx = out_idx; - const int x = idx % crop_width; - idx /= crop_width; - const int y = idx % crop_height; - idx /= crop_height; - const int d = idx % depth; - const int b = idx / depth; - - const float y1 = boxes_ptr[b * 4]; - const float x1 = boxes_ptr[b * 4 + 1]; - const float y2 = boxes_ptr[b * 4 + 2]; - const float x2 = boxes_ptr[b * 4 + 3]; - - const int b_in = box_ind_ptr[b]; - if (b_in < 0 || b_in >= batch) - { - continue; - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height - 1) / (crop_height - 1) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) : 0; - - const float in_y = (crop_height > 1) - ? y1 * (image_height - 1) + y * height_scale - : 0.5 * (y1 + y2) * (image_height - 1); - if (in_y < 0 || in_y > image_height - 1) - { - continue; - } - - const float in_x = (crop_width > 1) - ? x1 * (image_width - 1) + x * width_scale - : 0.5 * (x1 + x2) * (image_width - 1); - if (in_x < 0 || in_x > image_width - 1) - { - continue; - } - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - float *pimage = grads_image_ptr + (b_in * depth + d) * image_height * image_width; - const float dtop = (1 - y_lerp) * grads_ptr[out_idx]; - atomicAdd( - pimage + top_y_index * image_width + left_x_index, - (1 - x_lerp) * dtop - ); - atomicAdd( - pimage + top_y_index * image_width + right_x_index, - x_lerp * dtop - ); - - const float dbottom = y_lerp * grads_ptr[out_idx]; - atomicAdd( - pimage + bottom_y_index * image_width + left_x_index, - (1 - x_lerp) * dbottom - ); - atomicAdd( - pimage + bottom_y_index * image_width + right_x_index, - x_lerp * dbottom - ); - } -} - - -void CropAndResizeLaucher( - const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float extrapolation_value, float *crops_ptr, cudaStream_t stream) -{ - const int total_count = num_boxes * crop_height * crop_width * depth; - const int thread_per_block = 1024; - const int block_count = (total_count + thread_per_block - 1) / thread_per_block; - cudaError_t err; - - if (total_count > 0) - { - CropAndResizeKernel<<>>( - total_count, image_ptr, boxes_ptr, - box_ind_ptr, num_boxes, batch, image_height, image_width, - crop_height, crop_width, depth, extrapolation_value, crops_ptr); - - err = cudaGetLastError(); - if (cudaSuccess != err) - { - fprintf(stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } - } -} - - -void CropAndResizeBackpropImageLaucher( - const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float *grads_image_ptr, cudaStream_t stream) -{ - const int total_count = num_boxes * crop_height * crop_width * depth; - const int thread_per_block = 1024; - const int block_count = (total_count + thread_per_block - 1) / thread_per_block; - cudaError_t err; - - if (total_count > 0) - { - CropAndResizeBackpropImageKernel<<>>( - total_count, grads_ptr, boxes_ptr, - box_ind_ptr, num_boxes, batch, image_height, image_width, - crop_height, crop_width, depth, grads_image_ptr); - - err = cudaGetLastError(); - if (cudaSuccess != err) - { - fprintf(stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } - } -} \ No newline at end of file diff --git a/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.cu b/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.cu deleted file mode 100644 index 0702551..0000000 --- a/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.cu +++ /dev/null @@ -1,250 +0,0 @@ -#include -#include -#include "crop_and_resize_kernel.h" - -#define CUDA_1D_KERNEL_LOOP(i, n) \ -for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \ - i += blockDim.x * gridDim.x) - - -__global__ -void CropAndResizeKernel( - const int nthreads, const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float extrapolation_value, float *crops_ptr) -{ - CUDA_1D_KERNEL_LOOP(out_idx, nthreads) - { - // NHWC: out_idx = d + depth * (w + crop_width * (h + crop_height * b)) - // NCHW: out_idx = w + crop_width * (h + crop_height * (d + depth * b)) - int idx = out_idx; - //printf("start %i \n", idx); - const int x = idx % crop_width; - idx /= crop_width; - const int y = idx % crop_height; - idx /= crop_height; - const int d = idx % depth; - const int b = idx / depth; - - const float y1 = boxes_ptr[b * 4]; - const float x1 = boxes_ptr[b * 4 + 1]; - const float y2 = boxes_ptr[b * 4 + 2]; - const float x2 = boxes_ptr[b * 4 + 3]; - - const int b_in = box_ind_ptr[b]; - if (b_in < 0 || b_in >= batch) - { - continue; - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height) / (crop_height) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width) / (crop_width) : 0; - - - float tmp_in_y = (crop_height > 1) - ? y1 * (image_height ) + y * height_scale + height_scale/2 - 0.5 - : 0.5 * (y1 + y2) * (image_height); - - if (tmp_in_y > image_height - 1) - { - tmp_in_y = image_height - 1; - } - if (tmp_in_y < 0) - { - tmp_in_y = 0; - } - const float in_y = tmp_in_y; - - float tmp_in_x = (crop_width > 1) - ? x1 * (image_width ) + x * width_scale + width_scale/2 - 0.5 - : 0.5 * (x1 + x2) * (image_width ); - - if (tmp_in_x > image_width - 1) - { - tmp_in_x = image_width - 1; - } - if (tmp_in_x < 0) - { - tmp_in_x= 0; - } - const float in_x = tmp_in_x; - - //printf("height_scale %f \n", height_scale); - //printf("width_scale %f \n", width_scale); - //printf("in_x %f \n", in_x); - //printf("in_y %f \n", in_y); - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - const float *pimage = image_ptr + (b_in * depth + d) * image_height * image_width; - const float top_left = pimage[top_y_index * image_width + left_x_index]; - const float top_right = pimage[top_y_index * image_width + right_x_index]; - const float bottom_left = pimage[bottom_y_index * image_width + left_x_index]; - const float bottom_right = pimage[bottom_y_index * image_width + right_x_index]; - - const float top = top_left + (top_right - top_left) * x_lerp; - const float bottom = bottom_left + (bottom_right - bottom_left) * x_lerp; - crops_ptr[out_idx] = top + (bottom - top) * y_lerp; - } -} - -__global__ -void CropAndResizeBackpropImageKernel( - const int nthreads, const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float *grads_image_ptr) -{ - CUDA_1D_KERNEL_LOOP(out_idx, nthreads) - { - // NHWC: out_idx = d + depth * (w + crop_width * (h + crop_height * b)) - // NCHW: out_idx = w + crop_width * (h + crop_height * (d + depth * b)) - int idx = out_idx; - const int x = idx % crop_width; - idx /= crop_width; - const int y = idx % crop_height; - idx /= crop_height; - const int d = idx % depth; - const int b = idx / depth; - - const float y1 = boxes_ptr[b * 4]; - const float x1 = boxes_ptr[b * 4 + 1]; - const float y2 = boxes_ptr[b * 4 + 2]; - const float x2 = boxes_ptr[b * 4 + 3]; - - const int b_in = box_ind_ptr[b]; - if (b_in < 0 || b_in >= batch) - { - continue; - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height ) / (crop_height ) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width ) / (crop_width ) : 0; - - float tmp_in_y = (crop_height > 1) - ? y1 * (image_height ) + y * height_scale + height_scale/2 - 0.5 - : 0.5 * (y1 + y2) * (image_height); - - if (tmp_in_y > image_height - 1) - { - tmp_in_y = image_height - 1; - } - if (tmp_in_y < 0) - { - tmp_in_y = 0; - } - const float in_y = tmp_in_y; - - float tmp_in_x = (crop_width > 1) - ? x1 * (image_width ) + x * width_scale + width_scale/2 - 0.5 - : 0.5 * (x1 + x2) * (image_width ); - - if (tmp_in_x > image_width - 1) - { - tmp_in_x = image_width - 1; - } - if (tmp_in_x < 0) - { - tmp_in_x= 0; - } - const float in_x = tmp_in_x; - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - float *pimage = grads_image_ptr + (b_in * depth + d) * image_height * image_width; - const float dtop = (1 - y_lerp) * grads_ptr[out_idx]; - atomicAdd( - pimage + top_y_index * image_width + left_x_index, - (1 - x_lerp) * dtop - ); - atomicAdd( - pimage + top_y_index * image_width + right_x_index, - x_lerp * dtop - ); - - const float dbottom = y_lerp * grads_ptr[out_idx]; - atomicAdd( - pimage + bottom_y_index * image_width + left_x_index, - (1 - x_lerp) * dbottom - ); - atomicAdd( - pimage + bottom_y_index * image_width + right_x_index, - x_lerp * dbottom - ); - } -} - - -void CropAndResizeLaucher( - const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float extrapolation_value, float *crops_ptr, cudaStream_t stream) -{ - const int total_count = num_boxes * crop_height * crop_width * depth; - const int thread_per_block = 1024; - const int block_count = (total_count + thread_per_block - 1) / thread_per_block; - cudaError_t err; - - if (total_count > 0) - { - CropAndResizeKernel<<>>( - total_count, image_ptr, boxes_ptr, - box_ind_ptr, num_boxes, batch, image_height, image_width, - crop_height, crop_width, depth, extrapolation_value, crops_ptr); - - err = cudaGetLastError(); - if (cudaSuccess != err) - { - fprintf(stderr, "cudaCheckError in Roi Align () failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } - } -} - - -void CropAndResizeBackpropImageLaucher( - const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float *grads_image_ptr, cudaStream_t stream) -{ - const int total_count = num_boxes * crop_height * crop_width * depth; - const int thread_per_block = 1024; - const int block_count = (total_count + thread_per_block - 1) / thread_per_block; - cudaError_t err; - - if (total_count > 0) - { - CropAndResizeBackpropImageKernel<<>>( - total_count, grads_ptr, boxes_ptr, - box_ind_ptr, num_boxes, batch, image_height, image_width, - crop_height, crop_width, depth, grads_image_ptr); - - err = cudaGetLastError(); - if (cudaSuccess != err) - { - fprintf(stderr, "cudaCheckError() failed in Roi Align : %s\n", cudaGetErrorString(err)); - exit(-1); - } - } -} \ No newline at end of file diff --git a/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.cu.o b/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.cu.o deleted file mode 100644 index 2f1a1b9..0000000 Binary files a/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.cu.o and /dev/null differ diff --git a/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.h b/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.h deleted file mode 100644 index 893aee1..0000000 --- a/cuda_functions/roi_align_2D/roi_align/src/cuda/crop_and_resize_kernel.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef _CropAndResize_Kernel -#define _CropAndResize_Kernel - -#ifdef __cplusplus -extern "C" { -#endif - -void CropAndResizeLaucher( - const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float extrapolation_value, float *crops_ptr, cudaStream_t stream); - -void CropAndResizeBackpropImageLaucher( - const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float *grads_image_ptr, cudaStream_t stream); - -#ifdef __cplusplus -} -#endif - -#endif \ No newline at end of file diff --git a/cuda_functions/roi_align_2D/roi_align/src/cuda/fix.cu b/cuda_functions/roi_align_2D/roi_align/src/cuda/fix.cu deleted file mode 100644 index 6eea4a8..0000000 --- a/cuda_functions/roi_align_2D/roi_align/src/cuda/fix.cu +++ /dev/null @@ -1,243 +0,0 @@ -#include -#include -#include "crop_and_resize_kernel.h" - -#define CUDA_1D_KERNEL_LOOP(i, n) \ -for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \ - i += blockDim.x * gridDim.x) - - -__global__ -void CropAndResizeKernel( - const int nthreads, const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float extrapolation_value, float *crops_ptr) -{ - CUDA_1D_KERNEL_LOOP(out_idx, nthreads) - { - // NHWC: out_idx = d + depth * (w + crop_width * (h + crop_height * b)) - // NCHW: out_idx = w + crop_width * (h + crop_height * (d + depth * b)) - int idx = out_idx; - const int x = idx % crop_width; - idx /= crop_width; - const int y = idx % crop_height; - idx /= crop_height; - const int d = idx % depth; - const int b = idx / depth; - - const float y1 = boxes_ptr[b * 4]; - const float x1 = boxes_ptr[b * 4 + 1]; - const float y2 = boxes_ptr[b * 4 + 2]; - const float x2 = boxes_ptr[b * 4 + 3]; - - // printf("INIT CUDA SCRIPT %f \n", idx); - - const int b_in = box_ind_ptr[b]; - if (b_in < 0 || b_in >= batch) - { - continue; - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height ) / (crop_height ) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width) / (crop_width ) : 0; - - const float in_y = (crop_height > 1) - ? y1 * (image_height ) + y * height_scale + height_scale/2 - 0.5 - : 0.5 * (y1 + y2) * (image_height ); - if (in_y < 0 || in_y > image_height ) - { - crops_ptr[out_idx] = extrapolation_value; - continue; - } - - const float in_x = (crop_width > 1) - ? x1 * (image_width ) + x * width_scale + width_scale/2 - 0.5 - : 0.5 * (x1 + x2) * (image_width ); - if (in_x < 0 || in_x > image_width ) - { - crops_ptr[out_idx] = extrapolation_value; - continue; - } - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - const float *pimage = image_ptr + (b_in * depth + d) * image_height * image_width; - const float top_left = pimage[top_y_index * image_width + left_x_index]; - const float top_right = pimage[top_y_index * image_width + right_x_index]; - const float bottom_left = pimage[bottom_y_index * image_width + left_x_index]; - const float bottom_right = pimage[bottom_y_index * image_width + right_x_index]; - // if (top_left == 0){ - // const float top = top_right} - // elif (top_right == 0){ - // const float top = top_left} - // else{ - const float top = top_left + (top_right - top_left) * x_lerp; - //} - - //if (bottom_left == 0){ - // const float bottom = bottom_right} - // elif (bottom_right == 0){ - // const float bottom = bottom_left} - // else{ - const float bottom = bottom_left + (bottom_right - bottom_left) * x_lerp; - //} - - //if (top == 0){ - // crops_ptr[out_idx] = bottom } - // elif (bottom == 0){ - // crops_ptr[out_idx] = top - //} - // else{ - crops_ptr[out_idx] = top + (bottom - top) * y_lerp; - //} - } -} - -__global__ -void CropAndResizeBackpropImageKernel( - const int nthreads, const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float *grads_image_ptr) -{ - CUDA_1D_KERNEL_LOOP(out_idx, nthreads) - { - // NHWC: out_idx = d + depth * (w + crop_width * (h + crop_height * b)) - // NCHW: out_idx = w + crop_width * (h + crop_height * (d + depth * b)) - int idx = out_idx; - const int x = idx % crop_width; - idx /= crop_width; - const int y = idx % crop_height; - idx /= crop_height; - const int d = idx % depth; - const int b = idx / depth; - - const float y1 = boxes_ptr[b * 4]; - const float x1 = boxes_ptr[b * 4 + 1]; - const float y2 = boxes_ptr[b * 4 + 2]; - const float x2 = boxes_ptr[b * 4 + 3]; - - const int b_in = box_ind_ptr[b]; - if (b_in < 0 || b_in >= batch) - { - continue; - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height ) / (crop_height ) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width ) / (crop_width ) : 0; - - const float in_y = (crop_height > 1) - ? y1 * (image_height ) + y * height_scale + height_scale/2 - 0.5 - : 0.5 * (y1 + y2) * (image_height ); - if (in_y < 0 || in_y > image_height ) - { - continue; - } - - const float in_x = (crop_width > 1) - ? x1 * (image_width ) + x * width_scale + width_scale/2 - 0.5 - : 0.5 * (x1 + x2) * (image_width ); - if (in_x < 0 || in_x > image_width ) - { - continue; - } - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - float *pimage = grads_image_ptr + (b_in * depth + d) * image_height * image_width; - const float dtop = (1 - y_lerp) * grads_ptr[out_idx]; - atomicAdd( - pimage + top_y_index * image_width + left_x_index, - (1 - x_lerp) * dtop - ); - atomicAdd( - pimage + top_y_index * image_width + right_x_index, - x_lerp * dtop - ); - - const float dbottom = y_lerp * grads_ptr[out_idx]; - atomicAdd( - pimage + bottom_y_index * image_width + left_x_index, - (1 - x_lerp) * dbottom - ); - atomicAdd( - pimage + bottom_y_index * image_width + right_x_index, - x_lerp * dbottom - ); - } -} - - -void CropAndResizeLaucher( - const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float extrapolation_value, float *crops_ptr, cudaStream_t stream) -{ - const int total_count = num_boxes * crop_height * crop_width * depth; - const int thread_per_block = 1024; - const int block_count = (total_count + thread_per_block - 1) / thread_per_block; - cudaError_t err; - - if (total_count > 0) - { - CropAndResizeKernel<<>>( - total_count, image_ptr, boxes_ptr, - box_ind_ptr, num_boxes, batch, image_height, image_width, - crop_height, crop_width, depth, extrapolation_value, crops_ptr); - - err = cudaGetLastError(); - if (cudaSuccess != err) - { - fprintf(stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } - } -} - - -void CropAndResizeBackpropImageLaucher( - const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int crop_height, int crop_width, int depth, - float *grads_image_ptr, cudaStream_t stream) -{ - const int total_count = num_boxes * crop_height * crop_width * depth; - const int thread_per_block = 1024; - const int block_count = (total_count + thread_per_block - 1) / thread_per_block; - cudaError_t err; - - if (total_count > 0) - { - CropAndResizeBackpropImageKernel<<>>( - total_count, grads_ptr, boxes_ptr, - box_ind_ptr, num_boxes, batch, image_height, image_width, - crop_height, crop_width, depth, grads_image_ptr); - - err = cudaGetLastError(); - if (cudaSuccess != err) - { - fprintf(stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } - } -} \ No newline at end of file diff --git a/cuda_functions/roi_align_3D/__init__.py b/cuda_functions/roi_align_3D/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/roi_align_3D/roi_align/__init__.py b/cuda_functions/roi_align_3D/roi_align/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/roi_align_3D/roi_align/_ext/__init__.py b/cuda_functions/roi_align_3D/roi_align/_ext/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/._crop_and_resize.so.swp b/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/._crop_and_resize.so.swp deleted file mode 100644 index 3db0ea4..0000000 Binary files a/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/._crop_and_resize.so.swp and /dev/null differ diff --git a/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/__init__.py b/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/__init__.py deleted file mode 100644 index 4486c09..0000000 --- a/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ - -from torch.utils.ffi import _wrap_function -from ._crop_and_resize import lib as _lib, ffi as _ffi - -__all__ = [] -def _import_symbols(locals): - for symbol in dir(_lib): - fn = getattr(_lib, symbol) - if callable(fn): - locals[symbol] = _wrap_function(fn, _ffi) - else: - locals[symbol] = fn - __all__.append(symbol) - -_import_symbols(locals()) diff --git a/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/_crop_and_resize.so b/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/_crop_and_resize.so deleted file mode 100755 index 81dc147..0000000 Binary files a/cuda_functions/roi_align_3D/roi_align/_ext/crop_and_resize/_crop_and_resize.so and /dev/null differ diff --git a/cuda_functions/roi_align_3D/roi_align/build.py b/cuda_functions/roi_align_3D/roi_align/build.py deleted file mode 100755 index 3798d82..0000000 --- a/cuda_functions/roi_align_3D/roi_align/build.py +++ /dev/null @@ -1,40 +0,0 @@ -import os -import torch -from torch.utils.ffi import create_extension - - -sources = ['src/crop_and_resize.c'] -headers = ['src/crop_and_resize.h'] -defines = [] -with_cuda = False - -extra_objects = [] -if torch.cuda.is_available(): - print('Including CUDA code.') - sources += ['src/crop_and_resize_gpu.c'] - headers += ['src/crop_and_resize_gpu.h'] - defines += [('WITH_CUDA', None)] - extra_objects += ['src/cuda/crop_and_resize_kernel.cu.o'] - with_cuda = True - -extra_compile_args = ['-fopenmp', '-std=c99'] - -this_file = os.path.dirname(os.path.realpath(__file__)) -print(this_file) -sources = [os.path.join(this_file, fname) for fname in sources] -headers = [os.path.join(this_file, fname) for fname in headers] -extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] - -ffi = create_extension( - '_ext.crop_and_resize', - headers=headers, - sources=sources, - define_macros=defines, - relative_to=__file__, - with_cuda=with_cuda, - extra_objects=extra_objects, - extra_compile_args=extra_compile_args -) - -if __name__ == '__main__': - ffi.build() diff --git a/cuda_functions/roi_align_3D/roi_align/crop_and_resize.py b/cuda_functions/roi_align_3D/roi_align/crop_and_resize.py deleted file mode 100755 index cff4e90..0000000 --- a/cuda_functions/roi_align_3D/roi_align/crop_and_resize.py +++ /dev/null @@ -1,69 +0,0 @@ -import math -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.autograd import Function - -from ._ext import crop_and_resize as _backend - - -class CropAndResizeFunction(Function): - - def __init__(self, crop_height, crop_width, crop_zdepth, extrapolation_value=0): - self.crop_height = crop_height - self.crop_width = crop_width - self.crop_zdepth = crop_zdepth - self.extrapolation_value = extrapolation_value - - def forward(self, image, boxes, box_ind): - crops = torch.zeros_like(image) - - if image.is_cuda: - _backend.crop_and_resize_gpu_forward( - image, boxes, box_ind, - self.extrapolation_value, self.crop_height, self.crop_width, self.crop_zdepth, crops) - else: - _backend.crop_and_resize_forward( - image, boxes, box_ind, - self.extrapolation_value, self.crop_height, self.crop_width, self.crop_zdepth, crops) - - # save for backward - self.im_size = image.size() - self.save_for_backward(boxes, box_ind) - - return crops - - def backward(self, grad_outputs): - boxes, box_ind = self.saved_tensors - - grad_outputs = grad_outputs.contiguous() - grad_image = torch.zeros_like(grad_outputs).resize_(*self.im_size) - - if grad_outputs.is_cuda: - _backend.crop_and_resize_gpu_backward( - grad_outputs, boxes, box_ind, grad_image - ) - else: - _backend.crop_and_resize_backward( - grad_outputs, boxes, box_ind, grad_image - ) - - return grad_image, None, None - - -class CropAndResize(nn.Module): - """ - Crop and resize ported from tensorflow - See more details on https://www.tensorflow.org/api_docs/python/tf/image/crop_and_resize - """ - - def __init__(self, crop_height, crop_width, crop_zdepth, extrapolation_value=0): - super(CropAndResize, self).__init__() - - self.crop_height = crop_height - self.crop_width = crop_width - self.crop_zdepth = crop_zdepth - self.extrapolation_value = extrapolation_value - - def forward(self, image, boxes, box_ind): - return CropAndResizeFunction(self.crop_height, self.crop_width, self.crop_zdepth, self.extrapolation_value)(image, boxes, box_ind) diff --git a/cuda_functions/roi_align_3D/roi_align/roi_align.py b/cuda_functions/roi_align_3D/roi_align/roi_align.py deleted file mode 100644 index 6931539..0000000 --- a/cuda_functions/roi_align_3D/roi_align/roi_align.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -from torch import nn - -from .crop_and_resize import CropAndResizeFunction, CropAndResize - - -class RoIAlign(nn.Module): - - def __init__(self, crop_height, crop_width, extrapolation_value=0, transform_fpcoor=True): - super(RoIAlign, self).__init__() - - self.crop_height = crop_height - self.crop_width = crop_width - self.extrapolation_value = extrapolation_value - self.transform_fpcoor = transform_fpcoor - - def forward(self, featuremap, boxes, box_ind): - """ - RoIAlign based on crop_and_resize. - See more details on https://github.com/ppwwyyxx/tensorpack/blob/6d5ba6a970710eaaa14b89d24aace179eb8ee1af/examples/FasterRCNN/model.py#L301 - :param featuremap: NxCxHxW - :param boxes: Mx4 float box with (x1, y1, x2, y2) **without normalization** - :param box_ind: M - :return: MxCxoHxoW - """ - x1, y1, x2, y2 = torch.split(boxes, 1, dim=1) - image_height, image_width = featuremap.size()[2:4] - - if self.transform_fpcoor: - spacing_w = (x2 - x1) / float(self.crop_width) - spacing_h = (y2 - y1) / float(self.crop_height) - - nx0 = (x1 + spacing_w / 2 - 0.5) / float(image_width - 1) - ny0 = (y1 + spacing_h / 2 - 0.5) / float(image_height - 1) - nw = spacing_w * float(self.crop_width - 1) / float(image_width - 1) - nh = spacing_h * float(self.crop_height - 1) / float(image_height - 1) - - boxes = torch.cat((ny0, nx0, ny0 + nh, nx0 + nw), 1) - else: - x1 = x1 / float(image_width - 1) - x2 = x2 / float(image_width - 1) - y1 = y1 / float(image_height - 1) - y2 = y2 / float(image_height - 1) - boxes = torch.cat((y1, x1, y2, x2), 1) - - boxes = boxes.detach().contiguous() - box_ind = box_ind.detach() - return CropAndResizeFunction(self.crop_height, self.crop_width, self.extrapolation_value)(featuremap, boxes, box_ind) diff --git a/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize.c b/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize.c deleted file mode 100644 index a5ff973..0000000 --- a/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize.c +++ /dev/null @@ -1,269 +0,0 @@ -#include -#include -#include - - -void CropAndResizePerBox( - const float * image_data, - const int batch_size, - const int depth, - const int image_height, - const int image_width, - - const float * boxes_data, - const int * box_index_data, - const int start_box, - const int limit_box, - - float * corps_data, - const int crop_height, - const int crop_width, - const float extrapolation_value -) { - const int image_channel_elements = image_height * image_width; - const int image_elements = depth * image_channel_elements; - - const int channel_elements = crop_height * crop_width; - const int crop_elements = depth * channel_elements; - - int b; - #pragma omp parallel for - for (b = start_box; b < limit_box; ++b) { - const float * box = boxes_data + b * 4; - const float y1 = box[0]; - const float x1 = box[1]; - const float y2 = box[2]; - const float x2 = box[3]; - - const int b_in = box_index_data[b]; - if (b_in < 0 || b_in >= batch_size) { - printf("Error: batch_index %d out of range [0, %d)\n", b_in, batch_size); - exit(-1); - } - - const float height_scale = - (crop_height > 1) - ? (y2 - y1) * (image_height - 1) / (crop_height - 1) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) - : 0; - - for (int y = 0; y < crop_height; ++y) - { - const float in_y = (crop_height > 1) - ? y1 * (image_height - 1) + y * height_scale - : 0.5 * (y1 + y2) * (image_height - 1); - - if (in_y < 0 || in_y > image_height - 1) - { - for (int x = 0; x < crop_width; ++x) - { - for (int d = 0; d < depth; ++d) - { - // crops(b, y, x, d) = extrapolation_value; - corps_data[crop_elements * b + channel_elements * d + y * crop_width + x] = extrapolation_value; - } - } - continue; - } - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - for (int x = 0; x < crop_width; ++x) - { - const float in_x = (crop_width > 1) - ? x1 * (image_width - 1) + x * width_scale - : 0.5 * (x1 + x2) * (image_width - 1); - if (in_x < 0 || in_x > image_width - 1) - { - for (int d = 0; d < depth; ++d) - { - corps_data[crop_elements * b + channel_elements * d + y * crop_width + x] = extrapolation_value; - } - continue; - } - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - for (int d = 0; d < depth; ++d) - { - const float *pimage = image_data + b_in * image_elements + d * image_channel_elements; - - const float top_left = pimage[top_y_index * image_width + left_x_index]; - const float top_right = pimage[top_y_index * image_width + right_x_index]; - const float bottom_left = pimage[bottom_y_index * image_width + left_x_index]; - const float bottom_right = pimage[bottom_y_index * image_width + right_x_index]; - - const float top = top_left + (top_right - top_left) * x_lerp; - const float bottom = - bottom_left + (bottom_right - bottom_left) * x_lerp; - - corps_data[crop_elements * b + channel_elements * d + y * crop_width + x] = top + (bottom - top) * y_lerp; - } - } // end for x - } // end for y - } // end for b - -} - - -void crop_and_resize_forward( - THFloatTensor * image, - THFloatTensor * boxes, // [y1, x1, y2, x2] - THIntTensor * box_index, // range in [0, batch_size) - const float extrapolation_value, - const int crop_height, - const int crop_width, - THFloatTensor * crops -) { - //const int batch_size = image->size[0]; - //const int depth = image->size[1]; - //const int image_height = image->size[2]; - //const int image_width = image->size[3]; - - //const int num_boxes = boxes->size[0]; - - const int batch_size = THFloatTensor_size(image, 0); - const int depth = THFloatTensor_size(image, 1); - const int image_height = THFloatTensor_size(image, 2); - const int image_width = THFloatTensor_size(image, 3); - - const int num_boxes = THFloatTensor_size(boxes, 0); - - // init output space - THFloatTensor_resize4d(crops, num_boxes, depth, crop_height, crop_width); - THFloatTensor_zero(crops); - - // crop_and_resize for each box - CropAndResizePerBox( - THFloatTensor_data(image), - batch_size, - depth, - image_height, - image_width, - - THFloatTensor_data(boxes), - THIntTensor_data(box_index), - 0, - num_boxes, - - THFloatTensor_data(crops), - crop_height, - crop_width, - extrapolation_value - ); - -} - - -void crop_and_resize_backward( - THFloatTensor * grads, - THFloatTensor * boxes, // [y1, x1, y2, x2] - THIntTensor * box_index, // range in [0, batch_size) - THFloatTensor * grads_image // resize to [bsize, c, hc, wc] -) -{ - // shape - //const int batch_size = grads_image->size[0]; - //const int depth = grads_image->size[1]; - //const int image_height = grads_image->size[2]; - //const int image_width = grads_image->size[3]; - - //const int num_boxes = grads->size[0]; - //const int crop_height = grads->size[2]; - //const int crop_width = grads->size[3]; - - const int batch_size = THFloatTensor_size(grads_image, 0); - const int depth = THFloatTensor_size(grads_image, 1); - const int image_height = THFloatTensor_size(grads_image, 2); - const int image_width = THFloatTensor_size(grads_image, 3); - - const int num_boxes = THFloatTensor_size(grads, 0); - const int crop_height = THFloatTensor_size(grads,2); - const int crop_width = THFloatTensor_size(grads,3); - - - // n_elements - const int image_channel_elements = image_height * image_width; - const int image_elements = depth * image_channel_elements; - - const int channel_elements = crop_height * crop_width; - const int crop_elements = depth * channel_elements; - - // init output space - THFloatTensor_zero(grads_image); - - // data pointer - const float * grads_data = THFloatTensor_data(grads); - const float * boxes_data = THFloatTensor_data(boxes); - const int * box_index_data = THIntTensor_data(box_index); - float * grads_image_data = THFloatTensor_data(grads_image); - - for (int b = 0; b < num_boxes; ++b) { - const float * box = boxes_data + b * 4; - const float y1 = box[0]; - const float x1 = box[1]; - const float y2 = box[2]; - const float x2 = box[3]; - - const int b_in = box_index_data[b]; - if (b_in < 0 || b_in >= batch_size) { - printf("Error: batch_index %d out of range [0, %d)\n", b_in, batch_size); - exit(-1); - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height - 1) / (crop_height - 1) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) - : 0; - - for (int y = 0; y < crop_height; ++y) - { - const float in_y = (crop_height > 1) - ? y1 * (image_height - 1) + y * height_scale - : 0.5 * (y1 + y2) * (image_height - 1); - if (in_y < 0 || in_y > image_height - 1) - { - continue; - } - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - for (int x = 0; x < crop_width; ++x) - { - const float in_x = (crop_width > 1) - ? x1 * (image_width - 1) + x * width_scale - : 0.5 * (x1 + x2) * (image_width - 1); - if (in_x < 0 || in_x > image_width - 1) - { - continue; - } - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - for (int d = 0; d < depth; ++d) - { - float *pimage = grads_image_data + b_in * image_elements + d * image_channel_elements; - const float grad_val = grads_data[crop_elements * b + channel_elements * d + y * crop_width + x]; - - const float dtop = (1 - y_lerp) * grad_val; - pimage[top_y_index * image_width + left_x_index] += (1 - x_lerp) * dtop; - pimage[top_y_index * image_width + right_x_index] += x_lerp * dtop; - - const float dbottom = y_lerp * grad_val; - pimage[bottom_y_index * image_width + left_x_index] += (1 - x_lerp) * dbottom; - pimage[bottom_y_index * image_width + right_x_index] += x_lerp * dbottom; - } // end d - } // end x - } // end y - } // end b -} \ No newline at end of file diff --git a/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize.h b/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize.h deleted file mode 100644 index d494865..0000000 --- a/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize.h +++ /dev/null @@ -1,16 +0,0 @@ -void crop_and_resize_forward( - THFloatTensor * image, - THFloatTensor * boxes, // [y1, x1, y2, x2] - THIntTensor * box_index, // range in [0, batch_size) - const float extrapolation_value, - const int crop_height, - const int crop_width, - THFloatTensor * crops -); - -void crop_and_resize_backward( - THFloatTensor * grads, - THFloatTensor * boxes, // [y1, x1, y2, x2] - THIntTensor * box_index, // range in [0, batch_size) - THFloatTensor * grads_image // resize to [bsize, c, hc, wc] -); \ No newline at end of file diff --git a/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize_gpu.c b/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize_gpu.c deleted file mode 100644 index 8e07b3d..0000000 --- a/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize_gpu.c +++ /dev/null @@ -1,73 +0,0 @@ -#include -#include "cuda/crop_and_resize_kernel.h" - -extern THCState *state; - - -void crop_and_resize_gpu_forward( - THCudaTensor * image, - THCudaTensor * boxes, // [y1, x1, y2, x2] - THCudaIntTensor * box_index, // range in [0, batch_size) - const float extrapolation_value, - const int crop_height, - const int crop_width, - const int crop_zdepth, - THCudaTensor * crops -) { - const int batch_size = THCudaTensor_size(state, image, 0); - const int depth = THCudaTensor_size(state, image, 1); - const int image_height = THCudaTensor_size(state, image, 2); - const int image_width = THCudaTensor_size(state, image, 3); - const int image_zdepth = THCudaTensor_size(state, image, 4); - - const int num_boxes = THCudaTensor_size(state, boxes, 0); - - // init output space - THCudaTensor_resize5d(state, crops, num_boxes, depth, crop_height, crop_width, crop_zdepth); - THCudaTensor_zero(state, crops); - - cudaStream_t stream = THCState_getCurrentStream(state); - CropAndResizeLaucher( - THCudaTensor_data(state, image), - THCudaTensor_data(state, boxes), - THCudaIntTensor_data(state, box_index), - num_boxes, batch_size, image_height, image_width, image_zdepth, - crop_height, crop_width, crop_zdepth, depth, extrapolation_value, - THCudaTensor_data(state, crops), - stream - ); -} - - -void crop_and_resize_gpu_backward( - THCudaTensor * grads, - THCudaTensor * boxes, // [y1, x1, y2, x2] - THCudaIntTensor * box_index, // range in [0, batch_size) - THCudaTensor * grads_image // resize to [bsize, c, hc, wc] -) { - // shape - const int batch_size = THCudaTensor_size(state, grads_image, 0); - const int depth = THCudaTensor_size(state, grads_image, 1); - const int image_height = THCudaTensor_size(state, grads_image, 2); - const int image_width = THCudaTensor_size(state, grads_image, 3); - const int image_zdepth = THCudaTensor_size(state, grads_image, 4); - - const int num_boxes = THCudaTensor_size(state, grads, 0); - const int crop_height = THCudaTensor_size(state, grads, 2); - const int crop_width = THCudaTensor_size(state, grads, 3); - const int crop_zdepth = THCudaTensor_size(state, grads, 4); - - // init output space - THCudaTensor_zero(state, grads_image); - - cudaStream_t stream = THCState_getCurrentStream(state); - CropAndResizeBackpropImageLaucher( - THCudaTensor_data(state, grads), - THCudaTensor_data(state, boxes), - THCudaIntTensor_data(state, box_index), - num_boxes, batch_size, image_height, image_width, image_zdepth, - crop_height, crop_width, crop_zdepth, depth, - THCudaTensor_data(state, grads_image), - stream - ); -} \ No newline at end of file diff --git a/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize_gpu.h b/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize_gpu.h deleted file mode 100644 index dd2eb5a..0000000 --- a/cuda_functions/roi_align_3D/roi_align/src/crop_and_resize_gpu.h +++ /dev/null @@ -1,17 +0,0 @@ -void crop_and_resize_gpu_forward( - THCudaTensor * image, - THCudaTensor * boxes, // [y1, x1, y2, x2] - THCudaIntTensor * box_index, // range in [0, batch_size) - const float extrapolation_value, - const int crop_height, - const int crop_width, - const int crop_zdepth, - THCudaTensor * crops -); - -void crop_and_resize_gpu_backward( - THCudaTensor * grads, - THCudaTensor * boxes, // [y1, x1, y2, x2] - THCudaIntTensor * box_index, // range in [0, batch_size) - THCudaTensor * grads_image // resize to [bsize, c, hc, wc] -); \ No newline at end of file diff --git a/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.cu b/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.cu deleted file mode 100644 index e381dab..0000000 --- a/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.cu +++ /dev/null @@ -1,361 +0,0 @@ -#include -#include -#include "crop_and_resize_kernel.h" -#include - -#define CUDA_1D_KERNEL_LOOP(i, n) \ -for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \ - i += blockDim.x * gridDim.x) - - -__global__ -void CropAndResizeKernel( - const int nthreads, const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int image_zdepth, int crop_height, int crop_width, int crop_zdepth, int depth, - float extrapolation_value, float *crops_ptr) -{ - CUDA_1D_KERNEL_LOOP(out_idx, nthreads) // nthreads = total_count! - { - // NHWC: out_idx = d + depth * (w + crop_width * (h + crop_height * b)) position in out grid!!! - // NCHW: out_idx = w + crop_width * (h + crop_height * (d + depth * b)) NCYX yes seems like xy is exchanged! - // NCHWZ: out_idx = z + crop_zdepth * (w + crop_width * (h + crop_height * (d + depth * b))) z == last. - - int idx = out_idx; - - const int z = idx % crop_zdepth; - idx /= crop_zdepth; - const int x = idx % crop_width; - idx /= crop_width; - const int y = idx % crop_height; - idx /= crop_height; - - const int d = idx % depth; - const int b = idx / depth; // batch - - const float y1 = boxes_ptr[b * 6]; // b = batch -> 0 // normalized coords!! - const float x1 = boxes_ptr[b * 6 + 1]; - const float y2 = boxes_ptr[b * 6 + 2]; - const float x2 = boxes_ptr[b * 6 + 3]; - const float z1 = boxes_ptr[b * 6 + 4]; - const float z2 = boxes_ptr[b * 6 + 5]; - - const int b_in = box_ind_ptr[b]; // == 0 in my case. - if (b_in < 0 || b_in >= batch) - { - continue; - } - - // e.g. (0.4-0.3)*100 = 10 / 7 = 1.3 ratio proposal_size / crops_size. one cell in crops has size 1.3 in_pixel. - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height ) / (crop_height ) : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width ) / (crop_width ) : 0; - - const float zdepth_scale = - (crop_zdepth > 1) ? (z2 - z1) * (image_zdepth ) / (crop_zdepth ) : 0; - - - // e.g. 0.3*100 + 5 * 1.3 . Which floating coordinate is going into cell? - // e.g. y: 30 (lower bound prop) + 7.5 (current crop position * scale) - - - float tmp_in_y = (crop_height > 1) - ? y1 * (image_height ) + y * height_scale + height_scale/2 - 0.5 - : 0.5 * (y1 + y2) * (image_height); - - if (tmp_in_y > image_height - 1) - { - tmp_in_y = image_height - 1; - } - if (tmp_in_y < 0) - { - tmp_in_y = 0; - } - const float in_y = tmp_in_y; - - - float tmp_in_x = (crop_width > 1) - ? x1 * (image_width ) + x * width_scale + width_scale/2 - 0.5 - : 0.5 * (x1 + x2) * (image_width ); - - if (tmp_in_x > image_width - 1) - { - tmp_in_x = image_width - 1; - } - if (tmp_in_x < 0) - { - tmp_in_x= 0; - } - const float in_x = tmp_in_x; - - - float tmp_in_z = (crop_zdepth > 1) - ? z1 * (image_zdepth ) + z * zdepth_scale + zdepth_scale/2 - 0.5 - : 0.5 * (z1 + z2) * (image_zdepth); - - if (tmp_in_z > image_zdepth - 1) - { - tmp_in_z = image_zdepth - 1; - } - if (tmp_in_z < 0) - { - tmp_in_z= 0; - } - const float in_z = tmp_in_z; - - // this is just rounding of the floating coord of grid cell. The distances to nearest grid points are - // memorized (lerp) to be used for bilinear interpolation later. - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; // - - const int front_z_index = floorf(in_z); - const int back_z_index = ceilf(in_z); - const float z_lerp = in_z - front_z_index; - - - // address of image + going to the right feature map. - const float *pimage = image_ptr + (b_in * depth + d) * image_height * image_width * image_zdepth; - - // 1D address of corner points of in_coords to grid cell. - // NCHWZ: out_idx = z + crop_zdepth * (w + crop_width * (h + crop_height * (d + depth * b))) z == last. - const float top_left_front = pimage[front_z_index + image_zdepth * (left_x_index + image_width * top_y_index)]; - const float top_right_front = pimage[front_z_index + image_zdepth * (right_x_index + image_width * top_y_index)]; - const float bottom_left_front = pimage[front_z_index + image_zdepth * (left_x_index + image_width * bottom_y_index)]; - const float bottom_right_front = pimage[front_z_index + image_zdepth * (right_x_index + image_width * bottom_y_index)]; - const float top_left_back = pimage[back_z_index + image_zdepth * (left_x_index + image_width * top_y_index)]; - const float top_right_back = pimage[back_z_index + image_zdepth * (right_x_index + image_width * top_y_index)]; - const float bottom_left_back = pimage[back_z_index + image_zdepth * (left_x_index + image_width * bottom_y_index)]; - const float bottom_right_back = pimage[back_z_index + image_zdepth * (right_x_index + image_width * bottom_y_index)]; - - // Bilinear Interpolation!! These are pixel values now! lerp is the interpolation distance! - // No Maxpool, only one point is sampled! - const float top_front = top_left_front + (top_right_front - top_left_front) * x_lerp; - const float bottom_front = bottom_left_front + (bottom_right_front - bottom_left_front) * x_lerp; - const float top_back = top_left_back + (top_right_back - top_left_back) * x_lerp; - const float bottom_back = bottom_left_back + (bottom_right_back - bottom_left_back) * x_lerp; - - const float front = top_front + (bottom_front - top_front) * y_lerp; - const float back = top_back + (bottom_back - top_back) * y_lerp; - - crops_ptr[out_idx] = front + (back - front) * z_lerp; // assign interpolated value to Grid cell! - - - } -} - -__global__ -void CropAndResizeBackpropImageKernel( - const int nthreads, const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int image_zdepth, int crop_height, int crop_width, int crop_zdepth, int depth, - float *grads_image_ptr) -{ - CUDA_1D_KERNEL_LOOP(out_idx, nthreads) - { - // NHWC: out_idx = d + depth * (w + crop_width * (h + crop_height * b)) - // NCHW: out_idx = w + crop_width * (h + crop_height * (d + depth * b)) - // NCHWZ: out_idx = z + crop_zdepth * (w + crop_width * (h + crop_height * (d + depth * b))) z == last. - int idx = out_idx; - - const int z = idx % crop_zdepth; - idx /= crop_zdepth; - const int x = idx % crop_width; - idx /= crop_width; - const int y = idx % crop_height; - idx /= crop_height; - const int d = idx % depth; - const int b = idx / depth; - - const float y1 = boxes_ptr[b * 6]; // b = batch -> 0 // normalized coords!! - const float x1 = boxes_ptr[b * 6 + 1]; - const float y2 = boxes_ptr[b * 6 + 2]; - const float x2 = boxes_ptr[b * 6 + 3]; - const float z1 = boxes_ptr[b * 6 + 4]; - const float z2 = boxes_ptr[b * 6 + 5]; - - - const int b_in = box_ind_ptr[b]; - if (b_in < 0 || b_in >= batch) - { - continue; - } - - const float height_scale = - (crop_height > 1) ? (y2 - y1) * (image_height ) / (crop_height ) - : 0; - const float width_scale = - (crop_width > 1) ? (x2 - x1) * (image_width ) / (crop_width ) : 0; - - const float zdepth_scale = - (crop_zdepth > 1) ? (z2 - z1) * (image_zdepth ) / (crop_zdepth ) : 0; - - - float tmp_in_y = (crop_height > 1) - ? y1 * (image_height ) + y * height_scale + height_scale/2 - 0.5 - : 0.5 * (y1 + y2) * (image_height); - if (tmp_in_y > image_height - 1) - { - tmp_in_y = image_height - 1; - } - if (tmp_in_y < 0) - { - tmp_in_y = 0; - } - const float in_y = tmp_in_y; - - - float tmp_in_x = (crop_width > 1) - ? x1 * (image_width ) + x * width_scale + width_scale/2 - 0.5 - : 0.5 * (x1 + x2) * (image_width ); - if (tmp_in_x > image_width - 1) - { - tmp_in_x = image_width - 1; - } - if (tmp_in_x < 0) - { - tmp_in_x= 0; - } - const float in_x = tmp_in_x; - - - float tmp_in_z = (crop_zdepth > 1) - ? z1 * (image_zdepth ) + z * zdepth_scale + zdepth_scale/2 - 0.5 - : 0.5 * (z1 + z2) * (image_zdepth); - if (tmp_in_z > image_zdepth - 1) - { - tmp_in_z = image_zdepth - 1; - } - if (tmp_in_z < 0) - { - tmp_in_z= 0; - } - const float in_z = tmp_in_z; - - const int top_y_index = floorf(in_y); - const int bottom_y_index = ceilf(in_y); - const float y_lerp = in_y - top_y_index; - - const int left_x_index = floorf(in_x); - const int right_x_index = ceilf(in_x); - const float x_lerp = in_x - left_x_index; - - const int front_z_index = floorf(in_z); - const int back_z_index = ceilf(in_z); - const float z_lerp = in_z - front_z_index; - - float *pimage = grads_image_ptr + (b_in * depth + d) * image_height * image_width * image_zdepth; - - // top left front - atomicAdd( - pimage + front_z_index + image_zdepth * (left_x_index + image_width * top_y_index), - (1 - x_lerp) * (1 - z_lerp) * (1 - y_lerp) * grads_ptr[out_idx] // THIS IS BACKWARD INTERPOL. - ); - - // top left back - atomicAdd( - pimage + back_z_index + image_zdepth * (left_x_index + image_width * top_y_index), - (1 - x_lerp) * (z_lerp) * (1 - y_lerp) * grads_ptr[out_idx] // THIS IS BACKWARD INTERPOL. - ); - - // top right front - atomicAdd( - pimage + front_z_index + image_zdepth * (right_x_index + image_width * top_y_index), - (x_lerp) * (1 - z_lerp) * (1 - y_lerp) * grads_ptr[out_idx] // THIS IS backward INTERPOL. - ); - - // top right back - atomicAdd( - pimage + back_z_index + image_zdepth * (right_x_index + image_width * top_y_index), - (x_lerp) * (z_lerp) * (1 - y_lerp) * grads_ptr[out_idx] // THIS IS backward INTERPOL. - ); - - // bottom left front - atomicAdd( - pimage + front_z_index + image_zdepth * (left_x_index + image_width * bottom_y_index), - (1 - x_lerp) * (1 - z_lerp) * (y_lerp) * grads_ptr[out_idx] // THIS IS backward INTERPOL. - ); - - // bottom left back - atomicAdd( - pimage + back_z_index + image_zdepth * (left_x_index + image_width * bottom_y_index), - (1 - x_lerp) * (z_lerp) * (y_lerp) * grads_ptr[out_idx] // THIS IS backward INTERPOL. - ); - - // bottom right front - atomicAdd( - pimage + front_z_index + image_zdepth * (right_x_index + image_width * bottom_y_index), - (x_lerp) * (1 - z_lerp) * (y_lerp) * grads_ptr[out_idx] // THIS IS backward INTERPOL. - ); - - // bottom right back - atomicAdd( - pimage + back_z_index + image_zdepth * (right_x_index + image_width * bottom_y_index), - (x_lerp) * (z_lerp) * (y_lerp) * grads_ptr[out_idx] // THIS IS backward INTERPOL. - ); - - } -} - - - -void CropAndResizeLaucher( - const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int image_zdepth, int crop_height, int crop_width, int crop_zdepth, int depth, - float extrapolation_value, float *crops_ptr, cudaStream_t stream) -{ - const int total_count = num_boxes * crop_height * crop_width * crop_zdepth * depth; - const int thread_per_block = 1024; - const int block_count = (total_count + thread_per_block - 1) / thread_per_block; - cudaError_t err; - - if (total_count > 0) - { - CropAndResizeKernel<<>>( - total_count, image_ptr, boxes_ptr, - box_ind_ptr, num_boxes, batch, image_height, image_width, image_zdepth, - crop_height, crop_width, crop_zdepth, depth, extrapolation_value, crops_ptr); - - err = cudaGetLastError(); - if (cudaSuccess != err) - { - fprintf(stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString(err)); - exit(-1); - } - } -} - - -void CropAndResizeBackpropImageLaucher( - const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int image_zdepth, int crop_height, int crop_width, int crop_zdepth, int depth, - float *grads_image_ptr, cudaStream_t stream) -{ - const int total_count = num_boxes * crop_height * crop_width * crop_zdepth * depth; - const int thread_per_block = 1024; - const int block_count = (total_count + thread_per_block - 1) / thread_per_block; - cudaError_t err; - - if (total_count > 0) - { - CropAndResizeBackpropImageKernel<<>>( - total_count, grads_ptr, boxes_ptr, - box_ind_ptr, num_boxes, batch, image_height, image_width, image_zdepth, - crop_height, crop_width, crop_zdepth, depth, grads_image_ptr); - - err = cudaGetLastError(); - if (cudaSuccess != err) - { - fprintf(stderr, "cudaCheckError() failed in Roi Align : %s\n", cudaGetErrorString(err)); - exit(-1); - } - } -} \ No newline at end of file diff --git a/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.cu.o b/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.cu.o deleted file mode 100644 index d488598..0000000 Binary files a/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.cu.o and /dev/null differ diff --git a/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.h b/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.h deleted file mode 100644 index 9244582..0000000 --- a/cuda_functions/roi_align_3D/roi_align/src/cuda/crop_and_resize_kernel.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef _CropAndResize_Kernel -#define _CropAndResize_Kernel - -#ifdef __cplusplus -extern "C" { -#endif - -void CropAndResizeLaucher( - const float *image_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int image_zdepth, int crop_height, int crop_width, int crop_zdepth, int depth, - float extrapolation_value, float *crops_ptr, cudaStream_t stream); - -void CropAndResizeBackpropImageLaucher( - const float *grads_ptr, const float *boxes_ptr, - const int *box_ind_ptr, int num_boxes, int batch, int image_height, - int image_width, int image_zdepth, int crop_height, int crop_width, int crop_zdepth, int depth, - float *grads_image_ptr, cudaStream_t stream); - -#ifdef __cplusplus -} -#endif - -#endif \ No newline at end of file diff --git a/default_configs.py b/default_configs.py index 7abc68d..396e1d8 100644 --- a/default_configs.py +++ b/default_configs.py @@ -1,140 +1,140 @@ #!/usr/bin/env python # Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Default Configurations script. Avoids changing configs of all experiments if general settings are to be changed.""" import os class DefaultConfigs: def __init__(self, model, server_env=None, dim=2): - + self.server_env = server_env ######################### # I/O # ######################### self.model = model self.dim = dim # int [0 < dataset_size]. select n patients from dataset for prototyping. self.select_prototype_subset = None # some default paths. self.backbone_path = 'models/backbone.py' self.source_dir = os.path.dirname(os.path.realpath(__file__)) #current dir. self.input_df_name = 'info_df.pickle' self.model_path = 'models/{}.py'.format(self.model) if server_env: self.source_dir = '/home/jaegerp/code/mamma_code/medicaldetectiontoolkit' ######################### # Data Loader # ######################### #random seed for fold_generator and batch_generator. self.seed = 0 #number of threads for multithreaded batch generation. self.n_workers = 6 # if True, segmentation losses learn all categories, else only foreground vs. background. self.class_specific_seg_flag = False ######################### # Architecture # ######################### self.weight_decay = 0.0 # nonlinearity to be applied after convs with nonlinearity. one of 'relu' or 'leaky_relu' self.relu = 'relu' # if True initializes weights as specified in model script. else use default Pytorch init. self.custom_init = False # if True adds high-res decoder levels to feature pyramid: P1 + P0. (e.g. set to true in retina_unet configs) self.operate_stride1 = False ######################### # Schedule # ######################### # number of folds in cross validation. self.n_cv_splits = 5 # number of probabilistic samples in validation. self.n_probabilistic_samples = None ######################### # Testing / Plotting # ######################### # perform mirroring at test time. (only XY. Z not done to not blow up predictions times). self.test_aug = True # if True, test data lies in a separate folder and is not part of the cross validation. self.hold_out_test_set = False # if hold_out_test_set provided, ensemble predictions over models of all trained cv-folds. self.ensemble_folds = False # color specifications for all box_types in prediction_plot. self.box_color_palette = {'det': 'b', 'gt': 'r', 'neg_class': 'purple', 'prop': 'w', 'pos_class': 'g', 'pos_anchor': 'c', 'neg_anchor': 'c'} # scan over confidence score in evaluation to optimize it on the validation set. self.scan_det_thresh = False # plots roc-curves / prc-curves in evaluation. self.plot_stat_curves = False # evaluates average precision per image and averages over images. instead computing one ap over data set. self.per_patient_ap = False # threshold for clustering 2D box predictions to 3D Cubes. Overlap is computed in XY. self.merge_3D_iou = 0.1 # monitor any value from training. self.n_monitoring_figures = 1 # dict to assign specific plot_values to monitor_figures > 0. {1: ['class_loss'], 2: ['kl_loss', 'kl_sigmas']} self.assign_values_to_extra_figure = {} # save predictions to csv file in experiment dir. self.save_preds_to_csv = True # select a maximum number of patient cases to test. number or "all" for all self.max_test_patients = "all" ######################### # MRCNN # ######################### # if True, mask loss is not applied. used for data sets, where no pixel-wise annotations are provided. self.frcnn_mode = False # if True, unmolds masks in Mask R-CNN to full-res for plotting/monitoring. self.return_masks_in_val = False self.return_masks_in_test = False # needed if doing instance segmentation. evaluation not yet implemented. # add P6 to Feature Pyramid Network. self.sixth_pooling = False # for probabilistic detection self.n_latent_dims = 0 diff --git a/experiments/lidc_exp/configs.py b/experiments/lidc_exp/configs.py index f992f90..d0d640f 100644 --- a/experiments/lidc_exp/configs.py +++ b/experiments/lidc_exp/configs.py @@ -1,334 +1,334 @@ #!/usr/bin/env python # Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import sys import os sys.path.append(os.path.dirname(os.path.realpath(__file__))) import numpy as np from default_configs import DefaultConfigs class configs(DefaultConfigs): def __init__(self, server_env=None): ######################### # Preprocessing # ######################### self.root_dir = '/path/to/raw/data' self.raw_data_dir = '{}/data_nrrd'.format(self.root_dir) self.pp_dir = '{}/pp_norm'.format(self.root_dir) self.target_spacing = (0.7, 0.7, 1.25) ######################### # I/O # ######################### # one out of [2, 3]. dimension the model operates in. self.dim = 3 # one out of ['mrcnn', 'retina_net', 'retina_unet', 'detection_unet', 'ufrcnn']. - self.model = 'retina_unet' + self.model = 'mrcnn' DefaultConfigs.__init__(self, self.model, server_env, self.dim) # int [0 < dataset_size]. select n patients from dataset for prototyping. If None, all data is used. self.select_prototype_subset = None # path to preprocessed data. self.pp_name = 'lidc_preprocessed_for_G2' self.input_df_name = 'info_df.pickle' self.pp_data_path = '/mnt/HDD2TB/Documents/data/lidc/{}'.format(self.pp_name) self.pp_test_data_path = self.pp_data_path #change if test_data in separate folder. # settings for deployment in cloud. if server_env: # path to preprocessed data. self.pp_name = 'pp_fg_slices' self.crop_name = 'pp_fg_slices_packed' self.pp_data_path = '/path/to/preprocessed/data/{}/{}'.format(self.pp_name, self.crop_name) self.pp_test_data_path = self.pp_data_path self.select_prototype_subset = None ######################### # Data Loader # ######################### # select modalities from preprocessed data self.channels = [0] self.n_channels = len(self.channels) # patch_size to be used for training. pre_crop_size is the patch_size before data augmentation. self.pre_crop_size_2D = [300, 300] self.patch_size_2D = [288, 288] self.pre_crop_size_3D = [156, 156, 96] self.patch_size_3D = [128, 128, 64] self.patch_size = self.patch_size_2D if self.dim == 2 else self.patch_size_3D self.pre_crop_size = self.pre_crop_size_2D if self.dim == 2 else self.pre_crop_size_3D # ratio of free sampled batch elements before class balancing is triggered # (>0 to include "empty"/background patches.) self.batch_sample_slack = 0.2 # set 2D network to operate in 3D images. self.merge_2D_to_3D_preds = True # feed +/- n neighbouring slices into channel dimension. set to None for no context. self.n_3D_context = None if self.n_3D_context is not None and self.dim == 2: self.n_channels *= (self.n_3D_context * 2 + 1) ######################### # Architecture # ######################### self.start_filts = 48 if self.dim == 2 else 18 self.end_filts = self.start_filts * 4 if self.dim == 2 else self.start_filts * 2 self.res_architecture = 'resnet50' # 'resnet101' , 'resnet50' self.norm = None # one of None, 'instance_norm', 'batch_norm' self.weight_decay = 0 # one of 'xavier_uniform', 'xavier_normal', or 'kaiming_normal', None (=default = 'kaiming_uniform') self.weight_init = None ######################### # Schedule / Selection # ######################### self.num_epochs = 100 self.num_train_batches = 200 if self.dim == 2 else 200 self.batch_size = 20 if self.dim == 2 else 8 self.do_validation = True # decide whether to validate on entire patient volumes (like testing) or sampled patches (like training) # the former is morge accurate, while the latter is faster (depending on volume size) self.val_mode = 'val_sampling' # one of 'val_sampling' , 'val_patient' if self.val_mode == 'val_patient': self.max_val_patients = 50 # if 'None' iterates over entire val_set once. if self.val_mode == 'val_sampling': self.num_val_batches = 50 ######################### # Testing / Plotting # ######################### # set the top-n-epochs to be saved for temporal averaging in testing. self.save_n_models = 5 self.test_n_epochs = 5 # set a minimum epoch number for saving in case of instabilities in the first phase of training. self.min_save_thresh = 0 if self.dim == 2 else 0 self.report_score_level = ['patient', 'rois'] # choose list from 'patient', 'rois' self.class_dict = {1: 'benign', 2: 'malignant'} # 0 is background. self.patient_class_of_interest = 2 # patient metrics are only plotted for one class. self.ap_match_ious = [0.1] # list of ious to be evaluated for ap-scoring. self.model_selection_criteria = ['malignant_ap', 'benign_ap'] # criteria to average over for saving epochs. self.min_det_thresh = 0.1 # minimum confidence value to select predictions for evaluation. # threshold for clustering predictions together (wcs = weighted cluster scoring). # needs to be >= the expected overlap of predictions coming from one model (typically NMS threshold). # if too high, preds of the same object are separate clusters. self.wcs_iou = 1e-5 self.plot_prediction_histograms = True self.plot_stat_curves = False ######################### # Data Augmentation # ######################### self.da_kwargs={ 'do_elastic_deform': True, 'alpha':(0., 1500.), 'sigma':(30., 50.), 'do_rotation':True, 'angle_x': (0., 2 * np.pi), 'angle_y': (0., 0), 'angle_z': (0., 0), 'do_scale': True, 'scale':(0.8, 1.1), 'random_crop':False, 'rand_crop_dist': (self.patch_size[0] / 2. - 3, self.patch_size[1] / 2. - 3), 'border_mode_data': 'constant', 'border_cval_data': 0, 'order_data': 1 } if self.dim == 3: self.da_kwargs['do_elastic_deform'] = False self.da_kwargs['angle_x'] = (0, 0.0) self.da_kwargs['angle_y'] = (0, 0.0) #must be 0!! self.da_kwargs['angle_z'] = (0., 2 * np.pi) ######################### # Add model specifics # ######################### {'detection_unet': self.add_det_unet_configs, 'mrcnn': self.add_mrcnn_configs, 'ufrcnn': self.add_mrcnn_configs, 'retina_net': self.add_mrcnn_configs, 'retina_unet': self.add_mrcnn_configs, }[self.model]() def add_det_unet_configs(self): self.learning_rate = [1e-4] * self.num_epochs # aggregation from pixel perdiction to object scores (connected component). One of ['max', 'median'] self.aggregation_operation = 'max' # max number of roi candidates to identify per batch element and class. self.n_roi_candidates = 10 if self.dim == 2 else 30 # loss mode: either weighted cross entropy ('wce'), batch-wise dice loss ('dice), or the sum of both ('dice_wce') self.seg_loss_mode = 'dice_wce' # if <1, false positive predictions in foreground are penalized less. self.fp_dice_weight = 1 if self.dim == 2 else 1 self.wce_weights = [1, 1, 1] self.detection_min_confidence = self.min_det_thresh # if 'True', loss distinguishes all classes, else only foreground vs. background (class agnostic). self.class_specific_seg_flag = True self.num_seg_classes = 3 if self.class_specific_seg_flag else 2 self.head_classes = self.num_seg_classes def add_mrcnn_configs(self): # learning rate is a list with one entry per epoch. self.learning_rate = [1e-4] * self.num_epochs # disable the re-sampling of mask proposals to original size for speed-up. # since evaluation is detection-driven (box-matching) and not instance segmentation-driven (iou-matching), # mask-outputs are optional. self.return_masks_in_val = True self.return_masks_in_test = False # set number of proposal boxes to plot after each epoch. self.n_plot_rpn_props = 5 if self.dim == 2 else 30 # number of classes for head networks: n_foreground_classes + 1 (background) self.head_classes = 3 # seg_classes hier refers to the first stage classifier (RPN) self.num_seg_classes = 2 # foreground vs. background # feature map strides per pyramid level are inferred from architecture. self.backbone_strides = {'xy': [4, 8, 16, 32], 'z': [1, 2, 4, 8]} # anchor scales are chosen according to expected object sizes in data set. Default uses only one anchor scale # per pyramid level. (outer list are pyramid levels (corresponding to BACKBONE_STRIDES), inner list are scales per level.) self.rpn_anchor_scales = {'xy': [[8], [16], [32], [64]], 'z': [[2], [4], [8], [16]]} # choose which pyramid levels to extract features from: P2: 0, P3: 1, P4: 2, P5: 3. self.pyramid_levels = [0, 1, 2, 3] # number of feature maps in rpn. typically lowered in 3D to save gpu-memory. self.n_rpn_features = 512 if self.dim == 2 else 128 # anchor ratios and strides per position in feature maps. self.rpn_anchor_ratios = [0.5, 1, 2] self.rpn_anchor_stride = 1 # Threshold for first stage (RPN) non-maximum suppression (NMS): LOWER == HARDER SELECTION self.rpn_nms_threshold = 0.7 if self.dim == 2 else 0.7 # loss sampling settings. self.rpn_train_anchors_per_image = 6 #per batch element self.train_rois_per_image = 6 #per batch element self.roi_positive_ratio = 0.5 self.anchor_matching_iou = 0.7 # factor of top-k candidates to draw from per negative sample (stochastic-hard-example-mining). # poolsize to draw top-k candidates from will be shem_poolsize * n_negative_samples. self.shem_poolsize = 10 self.pool_size = (7, 7) if self.dim == 2 else (7, 7, 3) self.mask_pool_size = (14, 14) if self.dim == 2 else (14, 14, 5) self.mask_shape = (28, 28) if self.dim == 2 else (28, 28, 10) self.rpn_bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2]) self.bbox_std_dev = np.array([0.1, 0.1, 0.1, 0.2, 0.2, 0.2]) self.window = np.array([0, 0, self.patch_size[0], self.patch_size[1], 0, self.patch_size_3D[2]]) self.scale = np.array([self.patch_size[0], self.patch_size[1], self.patch_size[0], self.patch_size[1], self.patch_size_3D[2], self.patch_size_3D[2]]) if self.dim == 2: self.rpn_bbox_std_dev = self.rpn_bbox_std_dev[:4] self.bbox_std_dev = self.bbox_std_dev[:4] self.window = self.window[:4] self.scale = self.scale[:4] # pre-selection in proposal-layer (stage 1) for NMS-speedup. applied per batch element. self.pre_nms_limit = 3000 if self.dim == 2 else 6000 # n_proposals to be selected after NMS per batch element. too high numbers blow up memory if "detect_while_training" is True, # since proposals of the entire batch are forwarded through second stage in as one "batch". self.roi_chunk_size = 2500 if self.dim == 2 else 600 self.post_nms_rois_training = 500 if self.dim == 2 else 75 self.post_nms_rois_inference = 500 # Final selection of detections (refine_detections) self.model_max_instances_per_batch_element = 10 if self.dim == 2 else 30 # per batch element and class. self.detection_nms_threshold = 1e-5 # needs to be > 0, otherwise all predictions are one cluster. self.model_min_confidence = 0.1 if self.dim == 2: self.backbone_shapes = np.array( [[int(np.ceil(self.patch_size[0] / stride)), int(np.ceil(self.patch_size[1] / stride))] for stride in self.backbone_strides['xy']]) else: self.backbone_shapes = np.array( [[int(np.ceil(self.patch_size[0] / stride)), int(np.ceil(self.patch_size[1] / stride)), int(np.ceil(self.patch_size[2] / stride_z))] for stride, stride_z in zip(self.backbone_strides['xy'], self.backbone_strides['z'] )]) if self.model == 'ufrcnn': self.operate_stride1 = True self.class_specific_seg_flag = True self.num_seg_classes = 3 if self.class_specific_seg_flag else 2 self.frcnn_mode = True if self.model == 'retina_net' or self.model == 'retina_unet' or self.model == 'prob_detector': # implement extra anchor-scales according to retina-net publication. self.rpn_anchor_scales['xy'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in self.rpn_anchor_scales['xy']] self.rpn_anchor_scales['z'] = [[ii[0], ii[0] * (2 ** (1 / 3)), ii[0] * (2 ** (2 / 3))] for ii in self.rpn_anchor_scales['z']] self.n_anchors_per_pos = len(self.rpn_anchor_ratios) * 3 self.n_rpn_features = 256 if self.dim == 2 else 64 # pre-selection of detections for NMS-speedup. per entire batch. self.pre_nms_limit = 10000 if self.dim == 2 else 50000 # anchor matching iou is lower than in Mask R-CNN according to https://arxiv.org/abs/1708.02002 self.anchor_matching_iou = 0.5 # if 'True', seg loss distinguishes all classes, else only foreground vs. background (class agnostic). self.num_seg_classes = 3 if self.class_specific_seg_flag else 2 if self.model == 'retina_unet': self.operate_stride1 = True \ No newline at end of file diff --git a/experiments/lidc_exp/data_loader.py b/experiments/lidc_exp/data_loader.py index 2e64f34..0c8de96 100644 --- a/experiments/lidc_exp/data_loader.py +++ b/experiments/lidc_exp/data_loader.py @@ -1,461 +1,485 @@ #!/usr/bin/env python # Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== ''' Example Data Loader for the LIDC data set. This dataloader expects preprocessed data in .npy or .npz files per patient and a pandas dataframe in the same directory containing the meta-info e.g. file paths, labels, foregound slice-ids. ''' import numpy as np import os from collections import OrderedDict import pandas as pd import pickle import time import subprocess import utils.dataloader_utils as dutils # batch generator tools from https://github.com/MIC-DKFZ/batchgenerators from batchgenerators.dataloading.data_loader import SlimDataLoaderBase from batchgenerators.transforms.spatial_transforms import MirrorTransform as Mirror from batchgenerators.transforms.abstract_transforms import Compose from batchgenerators.dataloading.multi_threaded_augmenter import MultiThreadedAugmenter from batchgenerators.dataloading import SingleThreadedAugmenter from batchgenerators.transforms.spatial_transforms import SpatialTransform from batchgenerators.transforms.crop_and_pad_transforms import CenterCropTransform from batchgenerators.transforms.utility_transforms import ConvertSegToBoundingBoxCoordinates def get_train_generators(cf, logger): """ wrapper function for creating the training batch generator pipeline. returns the train/val generators. selects patients according to cv folds (generated by first run/fold of experiment): splits the data into n-folds, where 1 split is used for val, 1 split for testing and the rest for training. (inner loop test set) If cf.hold_out_test_set is True, adds the test split to the training data. """ all_data = load_dataset(cf, logger) all_pids_list = np.unique([v['pid'] for (k, v) in all_data.items()]) if not cf.created_fold_id_pickle: fg = dutils.fold_generator(seed=cf.seed, n_splits=cf.n_cv_splits, len_data=len(all_pids_list)).get_fold_names() with open(os.path.join(cf.exp_dir, 'fold_ids.pickle'), 'wb') as handle: pickle.dump(fg, handle) cf.created_fold_id_pickle = True else: with open(os.path.join(cf.exp_dir, 'fold_ids.pickle'), 'rb') as handle: fg = pickle.load(handle) train_ix, val_ix, test_ix, _ = fg[cf.fold] train_pids = [all_pids_list[ix] for ix in train_ix] val_pids = [all_pids_list[ix] for ix in val_ix] if cf.hold_out_test_set: train_pids += [all_pids_list[ix] for ix in test_ix] train_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in train_pids)} val_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in val_pids)} logger.info("data set loaded with: {} train / {} val / {} test patients".format(len(train_ix), len(val_ix), len(test_ix))) batch_gen = {} batch_gen['train'] = create_data_gen_pipeline(train_data, cf=cf, is_training=True) batch_gen['val_sampling'] = create_data_gen_pipeline(val_data, cf=cf, is_training=False) if cf.val_mode == 'val_patient': batch_gen['val_patient'] = PatientBatchIterator(val_data, cf=cf) batch_gen['n_val'] = len(val_ix) if cf.max_val_patients is None else min(len(val_ix), cf.max_val_patients) else: batch_gen['n_val'] = cf.num_val_batches return batch_gen def get_test_generator(cf, logger): """ wrapper function for creating the test batch generator pipeline. selects patients according to cv folds (generated by first run/fold of experiment) If cf.hold_out_test_set is True, gets the data from an external folder instead. """ if cf.hold_out_test_set: pp_name = cf.pp_test_name test_ix = None else: pp_name = None with open(os.path.join(cf.exp_dir, 'fold_ids.pickle'), 'rb') as handle: fold_list = pickle.load(handle) _, _, test_ix, _ = fold_list[cf.fold] # warnings.warn('WARNING: using validation set for testing!!!') test_data = load_dataset(cf, logger, test_ix, pp_data_path=cf.pp_test_data_path, pp_name=pp_name) logger.info("data set loaded with: {} test patients".format(len(test_ix))) batch_gen = {} batch_gen['test'] = PatientBatchIterator(test_data, cf=cf) batch_gen['n_test'] = len(test_ix) if cf.max_test_patients=="all" else \ min(cf.max_test_patients, len(test_ix)) return batch_gen def load_dataset(cf, logger, subset_ixs=None, pp_data_path=None, pp_name=None): """ loads the dataset. if deployed in cloud also copies and unpacks the data to the working directory. :param subset_ixs: subset indices to be loaded from the dataset. used e.g. for testing to only load the test folds. :return: data: dictionary with one entry per patient (in this case per patient-breast, since they are treated as individual images for training) each entry is a dictionary containing respective meta-info as well as paths to the preprocessed numpy arrays to be loaded during batch-generation """ if pp_data_path is None: pp_data_path = cf.pp_data_path if pp_name is None: pp_name = cf.pp_name if cf.server_env: copy_data = True target_dir = os.path.join('/ssd', cf.slurm_job_id, pp_name, cf.crop_name) if not os.path.exists(target_dir): cf.data_source_dir = pp_data_path os.makedirs(target_dir) subprocess.call('rsync -av {} {}'.format( os.path.join(cf.data_source_dir, cf.input_df_name), os.path.join(target_dir, cf.input_df_name)), shell=True) logger.info('created target dir and info df at {}'.format(os.path.join(target_dir, cf.input_df_name))) elif subset_ixs is None: copy_data = False pp_data_path = target_dir p_df = pd.read_pickle(os.path.join(pp_data_path, cf.input_df_name)) if cf.select_prototype_subset is not None: prototype_pids = p_df.pid.tolist()[:cf.select_prototype_subset] p_df = p_df[p_df.pid.isin(prototype_pids)] logger.warning('WARNING: using prototyping data subset!!!') if subset_ixs is not None: subset_pids = [np.unique(p_df.pid.tolist())[ix] for ix in subset_ixs] p_df = p_df[p_df.pid.isin(subset_pids)] logger.info('subset: selected {} instances from df'.format(len(p_df))) if cf.server_env: if copy_data: copy_and_unpack_data(logger, p_df.pid.tolist(), cf.fold_dir, cf.data_source_dir, target_dir) class_targets = p_df['class_target'].tolist() pids = p_df.pid.tolist() imgs = [os.path.join(pp_data_path, '{}_img.npy'.format(pid)) for pid in pids] segs = [os.path.join(pp_data_path,'{}_rois.npy'.format(pid)) for pid in pids] data = OrderedDict() for ix, pid in enumerate(pids): # for the experiment conducted here, malignancy scores are binarized: (benign: 1-2, malignant: 3-5) targets = [1 if ii >= 3 else 0 for ii in class_targets[ix]] data[pid] = {'data': imgs[ix], 'seg': segs[ix], 'pid': pid, 'class_target': targets} data[pid]['fg_slices'] = p_df['fg_slices'].tolist()[ix] return data def create_data_gen_pipeline(patient_data, cf, is_training=True): """ create mutli-threaded train/val/test batch generation and augmentation pipeline. :param patient_data: dictionary containing one dictionary per patient in the train/test subset. :param is_training: (optional) whether to perform data augmentation (training) or not (validation/testing) :return: multithreaded_generator """ # create instance of batch generator as first element in pipeline. data_gen = BatchGenerator(patient_data, batch_size=cf.batch_size, cf=cf) # add transformations to pipeline. my_transforms = [] if is_training: mirror_transform = Mirror(axes=np.arange(cf.dim)) my_transforms.append(mirror_transform) spatial_transform = SpatialTransform(patch_size=cf.patch_size[:cf.dim], patch_center_dist_from_border=cf.da_kwargs['rand_crop_dist'], do_elastic_deform=cf.da_kwargs['do_elastic_deform'], alpha=cf.da_kwargs['alpha'], sigma=cf.da_kwargs['sigma'], do_rotation=cf.da_kwargs['do_rotation'], angle_x=cf.da_kwargs['angle_x'], angle_y=cf.da_kwargs['angle_y'], angle_z=cf.da_kwargs['angle_z'], do_scale=cf.da_kwargs['do_scale'], scale=cf.da_kwargs['scale'], random_crop=cf.da_kwargs['random_crop']) my_transforms.append(spatial_transform) else: my_transforms.append(CenterCropTransform(crop_size=cf.patch_size[:cf.dim])) my_transforms.append(ConvertSegToBoundingBoxCoordinates(cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=cf.class_specific_seg_flag)) all_transforms = Compose(my_transforms) # multithreaded_generator = SingleThreadedAugmenter(data_gen, all_transforms) multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes=cf.n_workers, seeds=range(cf.n_workers)) return multithreaded_generator class BatchGenerator(SlimDataLoaderBase): """ creates the training/validation batch generator. Samples n_batch_size patients (draws a slice from each patient if 2D) from the data set while maintaining foreground-class balance. Returned patches are cropped/padded to pre_crop_size. Actual patch_size is obtained after data augmentation. :param data: data dictionary as provided by 'load_dataset'. :param batch_size: number of patients to sample for the batch :return dictionary containing the batch data (b, c, x, y, (z)) / seg (b, 1, x, y, (z)) / pids / class_target """ def __init__(self, data, batch_size, cf): super(BatchGenerator, self).__init__(data, batch_size) self.cf = cf self.crop_margin = np.array(self.cf.patch_size)/8. #min distance of ROI center to edge of cropped_patch. self.p_fg = 0.5 def generate_train_batch(self): batch_data, batch_segs, batch_pids, batch_targets, batch_patient_labels = [], [], [], [], [] class_targets_list = [v['class_target'] for (k, v) in self._data.items()] if self.cf.head_classes > 2: # samples patients towards equilibrium of foreground classes on a roi-level (after randomly sampling the ratio "batch_sample_slack). batch_ixs = dutils.get_class_balanced_patients( class_targets_list, self.batch_size, self.cf.head_classes - 1, slack_factor=self.cf.batch_sample_slack) else: batch_ixs = np.random.choice(len(class_targets_list), self.batch_size) patients = list(self._data.items()) for b in batch_ixs: patient = patients[b][1] data = np.transpose(np.load(patient['data'], mmap_mode='r'), axes=(1, 2, 0))[np.newaxis] # (c, y, x, z) seg = np.transpose(np.load(patient['seg'], mmap_mode='r'), axes=(1, 2, 0)) batch_pids.append(patient['pid']) batch_targets.append(patient['class_target']) if self.cf.dim == 2: # draw random slice from patient while oversampling slices containing foreground objects with p_fg. if len(patient['fg_slices']) > 0: fg_prob = self.p_fg / len(patient['fg_slices']) bg_prob = (1 - self.p_fg) / (data.shape[3] - len(patient['fg_slices'])) slices_prob = [fg_prob if ix in patient['fg_slices'] else bg_prob for ix in range(data.shape[3])] slice_id = np.random.choice(data.shape[3], p=slices_prob) else: slice_id = np.random.choice(data.shape[3]) # if set to not None, add neighbouring slices to each selected slice in channel dimension. if self.cf.n_3D_context is not None: padded_data = dutils.pad_nd_image(data[0], [(data.shape[-1] + (self.cf.n_3D_context*2))], mode='constant') padded_slice_id = slice_id + self.cf.n_3D_context data = (np.concatenate([padded_data[..., ii][np.newaxis] for ii in range( padded_slice_id - self.cf.n_3D_context, padded_slice_id + self.cf.n_3D_context + 1)], axis=0)) else: data = data[..., slice_id] seg = seg[..., slice_id] # pad data if smaller than pre_crop_size. if np.any([data.shape[dim + 1] < ps for dim, ps in enumerate(self.cf.pre_crop_size)]): new_shape = [np.max([data.shape[dim + 1], ps]) for dim, ps in enumerate(self.cf.pre_crop_size)] data = dutils.pad_nd_image(data, new_shape, mode='constant') seg = dutils.pad_nd_image(seg, new_shape, mode='constant') # crop patches of size pre_crop_size, while sampling patches containing foreground with p_fg. crop_dims = [dim for dim, ps in enumerate(self.cf.pre_crop_size) if data.shape[dim + 1] > ps] if len(crop_dims) > 0: fg_prob_sample = np.random.rand(1) # with p_fg: sample random pixel from random ROI and shift center by random value. if fg_prob_sample < self.p_fg and np.sum(seg) > 0: seg_ixs = np.argwhere(seg == np.random.choice(np.unique(seg)[1:], 1)) roi_anchor_pixel = seg_ixs[np.random.choice(seg_ixs.shape[0], 1)][0] assert seg[tuple(roi_anchor_pixel)] > 0 # sample the patch center coords. constrained by edges of images - pre_crop_size /2. And by # distance to the desired ROI < patch_size /2. # (here final patch size to account for center_crop after data augmentation). sample_seg_center = {} for ii in crop_dims: low = np.max((self.cf.pre_crop_size[ii]//2, roi_anchor_pixel[ii] - (self.cf.patch_size[ii]//2 - self.crop_margin[ii]))) high = np.min((data.shape[ii + 1] - self.cf.pre_crop_size[ii]//2, roi_anchor_pixel[ii] + (self.cf.patch_size[ii]//2 - self.crop_margin[ii]))) # happens if lesion on the edge of the image. dont care about roi anymore, # just make sure pre-crop is inside image. if low >= high: low = data.shape[ii + 1] // 2 - (data.shape[ii + 1] // 2 - self.cf.pre_crop_size[ii] // 2) high = data.shape[ii + 1] // 2 + (data.shape[ii + 1] // 2 - self.cf.pre_crop_size[ii] // 2) sample_seg_center[ii] = np.random.randint(low=low, high=high) else: # not guaranteed to be empty. probability of emptiness depends on the data. sample_seg_center = {ii: np.random.randint(low=self.cf.pre_crop_size[ii]//2, high=data.shape[ii + 1] - self.cf.pre_crop_size[ii]//2) for ii in crop_dims} for ii in crop_dims: min_crop = int(sample_seg_center[ii] - self.cf.pre_crop_size[ii] // 2) max_crop = int(sample_seg_center[ii] + self.cf.pre_crop_size[ii] // 2) data = np.take(data, indices=range(min_crop, max_crop), axis=ii + 1) seg = np.take(seg, indices=range(min_crop, max_crop), axis=ii) batch_data.append(data) batch_segs.append(seg[np.newaxis]) data = np.array(batch_data) seg = np.array(batch_segs).astype(np.uint8) class_target = np.array(batch_targets) return {'data': data, 'seg': seg, 'pid': batch_pids, 'class_target': class_target} class PatientBatchIterator(SlimDataLoaderBase): """ creates a test generator that iterates over entire given dataset returning 1 patient per batch. Can be used for monitoring if cf.val_mode = 'patient_val' for a monitoring closer to actualy evaluation (done in 3D), if willing to accept speed-loss during training. :return: out_batch: dictionary containing one patient with batch_size = n_3D_patches in 3D or batch_size = n_2D_patches in 2D . """ def __init__(self, data, cf): #threads in augmenter super(PatientBatchIterator, self).__init__(data, 0) self.cf = cf self.patient_ix = 0 self.dataset_pids = [v['pid'] for (k, v) in data.items()] self.patch_size = cf.patch_size if len(self.patch_size) == 2: self.patch_size = self.patch_size + [1] def generate_train_batch(self): pid = self.dataset_pids[self.patient_ix] patient = self._data[pid] data = np.transpose(np.load(patient['data'], mmap_mode='r'), axes=(1, 2, 0))[np.newaxis] # (c, y, x, z) seg = np.transpose(np.load(patient['seg'], mmap_mode='r'), axes=(1, 2, 0)) batch_class_targets = np.array([patient['class_target']]) # pad data if smaller than patch_size seen during training. if np.any([data.shape[dim + 1] < ps for dim, ps in enumerate(self.patch_size)]): new_shape = [data.shape[0]] + [np.max([data.shape[dim + 1], self.patch_size[dim]]) for dim, ps in enumerate(self.patch_size)] data = dutils.pad_nd_image(data, new_shape) # use 'return_slicer' to crop image back to original shape. seg = dutils.pad_nd_image(seg, new_shape) # get 3D targets for evaluation, even if network operates in 2D. 2D predictions will be merged to 3D in predictor. if self.cf.dim == 3 or self.cf.merge_2D_to_3D_preds: out_data = data[np.newaxis] out_seg = seg[np.newaxis, np.newaxis] out_targets = batch_class_targets batch_3D = {'data': out_data, 'seg': out_seg, 'class_target': out_targets, 'pid': pid} converter = ConvertSegToBoundingBoxCoordinates(dim=3, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag) batch_3D = converter(**batch_3D) batch_3D.update({'patient_bb_target': batch_3D['bb_target'], 'patient_roi_labels': batch_3D['roi_labels'], 'original_img_shape': out_data.shape}) if self.cf.dim == 2: out_data = np.transpose(data, axes=(3, 0, 1, 2)) # (z, c, x, y ) out_seg = np.transpose(seg, axes=(2, 0, 1))[:, np.newaxis] out_targets = np.array(np.repeat(batch_class_targets, out_data.shape[0], axis=0)) # if set to not None, add neighbouring slices to each selected slice in channel dimension. if self.cf.n_3D_context is not None: slice_range = range(self.cf.n_3D_context, out_data.shape[0] + self.cf.n_3D_context) out_data = np.pad(out_data, ((self.cf.n_3D_context, self.cf.n_3D_context), (0, 0), (0, 0), (0, 0)), 'constant', constant_values=0) out_data = np.array( [np.concatenate([out_data[ii] for ii in range( slice_id - self.cf.n_3D_context, slice_id + self.cf.n_3D_context + 1)], axis=0) for slice_id in slice_range]) batch_2D = {'data': out_data, 'seg': out_seg, 'class_target': out_targets, 'pid': pid} converter = ConvertSegToBoundingBoxCoordinates(dim=2, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag) batch_2D = converter(**batch_2D) if self.cf.merge_2D_to_3D_preds: batch_2D.update({'patient_bb_target': batch_3D['patient_bb_target'], 'patient_roi_labels': batch_3D['patient_roi_labels'], 'original_img_shape': out_data.shape}) else: batch_2D.update({'patient_bb_target': batch_2D['bb_target'], 'patient_roi_labels': batch_2D['roi_labels'], 'original_img_shape': out_data.shape}) out_batch = batch_3D if self.cf.dim == 3 else batch_2D patient_batch = out_batch # crop patient-volume to patches of patch_size used during training. stack patches up in batch dimension. # in this case, 2D is treated as a special case of 3D with patch_size[z] = 1. if np.any([data.shape[dim + 1] > self.patch_size[dim] for dim in range(3)]): patch_crop_coords_list = dutils.get_patch_crop_coords(data[0], self.patch_size) new_img_batch, new_seg_batch, new_class_targets_batch = [], [], [] for cix, c in enumerate(patch_crop_coords_list): seg_patch = seg[c[0]:c[1], c[2]: c[3], c[4]:c[5]] new_seg_batch.append(seg_patch) # if set to not None, add neighbouring slices to each selected slice in channel dimension. # correct patch_crop coordinates by added slices of 3D context. if self.cf.dim == 2 and self.cf.n_3D_context is not None: tmp_c_5 = c[5] + (self.cf.n_3D_context * 2) if cix == 0: data = np.pad(data, ((0, 0), (0, 0), (0, 0), (self.cf.n_3D_context, self.cf.n_3D_context)), 'constant', constant_values=0) else: tmp_c_5 = c[5] new_img_batch.append(data[:, c[0]:c[1], c[2]:c[3], c[4]:tmp_c_5]) data = np.array(new_img_batch) # (n_patches, c, x, y, z) seg = np.array(new_seg_batch)[:, np.newaxis] # (n_patches, 1, x, y, z) batch_class_targets = np.repeat(batch_class_targets, len(patch_crop_coords_list), axis=0) if self.cf.dim == 2: if self.cf.n_3D_context is not None: data = np.transpose(data[:, 0], axes=(0, 3, 1, 2)) else: # all patches have z dimension 1 (slices). discard dimension data = data[..., 0] seg = seg[..., 0] patch_batch = {'data': data, 'seg': seg, 'class_target': batch_class_targets, 'pid': pid} patch_batch['patch_crop_coords'] = np.array(patch_crop_coords_list) patch_batch['patient_bb_target'] = patient_batch['patient_bb_target'] patch_batch['patient_roi_labels'] = patient_batch['patient_roi_labels'] patch_batch['original_img_shape'] = patient_batch['original_img_shape'] converter = ConvertSegToBoundingBoxCoordinates(self.cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag) patch_batch = converter(**patch_batch) out_batch = patch_batch self.patient_ix += 1 if self.patient_ix == len(self.dataset_pids): self.patient_ix = 0 return out_batch def copy_and_unpack_data(logger, pids, fold_dir, source_dir, target_dir): start_time = time.time() with open(os.path.join(fold_dir, 'file_list.txt'), 'w') as handle: for pid in pids: handle.write('{}_img.npz\n'.format(pid)) handle.write('{}_rois.npz\n'.format(pid)) subprocess.call('rsync -av --files-from {} {} {}'.format(os.path.join(fold_dir, 'file_list.txt'), source_dir, target_dir), shell=True) dutils.unpack_dataset(target_dir) copied_files = os.listdir(target_dir) logger.info("copying and unpacking data set finsihed : {} files in target dir: {}. took {} sec".format( len(copied_files), target_dir, np.round(time.time() - start_time, 0))) + +if __name__=="__main__": + import utils.exp_utils as utils + from configs import configs + + total_stime = time.time() + + + cf = configs() + cf.created_fold_id_pickle = False + cf.exp_dir = "experiments/dev/" + cf.plot_dir = cf.exp_dir + "plots" + os.makedirs(cf.exp_dir, exist_ok=True) + cf.fold = 0 + logger = utils.get_logger(cf.exp_dir) + batch_gen = get_train_generators(cf, logger) + + train_batch = next(batch_gen["train"]) + + + mins, secs = divmod((time.time() - total_stime), 60) + h, mins = divmod(mins, 60) + t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs)) + print("{} total runtime: {}".format(os.path.split(__file__)[1], t)) \ No newline at end of file diff --git a/experiments/toy_exp/data_loader.py b/experiments/toy_exp/data_loader.py index 3a7062c..22e76f9 100644 --- a/experiments/toy_exp/data_loader.py +++ b/experiments/toy_exp/data_loader.py @@ -1,305 +1,305 @@ #!/usr/bin/env python # Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import numpy as np import os from collections import OrderedDict import pandas as pd import pickle import time import subprocess import utils.dataloader_utils as dutils # batch generator tools from https://github.com/MIC-DKFZ/batchgenerators from batchgenerators.dataloading.data_loader import SlimDataLoaderBase from batchgenerators.transforms.spatial_transforms import MirrorTransform as Mirror from batchgenerators.transforms.abstract_transforms import Compose from batchgenerators.dataloading.multi_threaded_augmenter import MultiThreadedAugmenter from batchgenerators.dataloading import SingleThreadedAugmenter from batchgenerators.transforms.spatial_transforms import SpatialTransform from batchgenerators.transforms.crop_and_pad_transforms import CenterCropTransform from batchgenerators.transforms.utility_transforms import ConvertSegToBoundingBoxCoordinates def get_train_generators(cf, logger): """ wrapper function for creating the training batch generator pipeline. returns the train/val generators. selects patients according to cv folds (generated by first run/fold of experiment): splits the data into n-folds, where 1 split is used for val, 1 split for testing and the rest for training. (inner loop test set) If cf.hold_out_test_set is True, adds the test split to the training data. """ all_data = load_dataset(cf, logger) all_pids_list = np.unique([v['pid'] for (k, v) in all_data.items()]) train_pids = all_pids_list[:cf.n_train_data] val_pids = all_pids_list[1000:1500] train_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in train_pids)} val_data = {k: v for (k, v) in all_data.items() if any(p == v['pid'] for p in val_pids)} logger.info("data set loaded with: {} train / {} val patients".format(len(train_pids), len(val_pids))) batch_gen = {} batch_gen['train'] = create_data_gen_pipeline(train_data, cf=cf, do_aug=False) batch_gen['val_sampling'] = create_data_gen_pipeline(val_data, cf=cf, do_aug=False) if cf.val_mode == 'val_patient': batch_gen['val_patient'] = PatientBatchIterator(val_data, cf=cf) batch_gen['n_val'] = len(val_pids) if cf.max_val_patients is None else min(len(val_pids), cf.max_val_patients) else: batch_gen['n_val'] = cf.num_val_batches return batch_gen def get_test_generator(cf, logger): """ wrapper function for creating the test batch generator pipeline. selects patients according to cv folds (generated by first run/fold of experiment) If cf.hold_out_test_set is True, gets the data from an external folder instead. """ if cf.hold_out_test_set: pp_name = cf.pp_test_name test_ix = None else: pp_name = None with open(os.path.join(cf.exp_dir, 'fold_ids.pickle'), 'rb') as handle: fold_list = pickle.load(handle) _, _, test_ix, _ = fold_list[cf.fold] # warnings.warn('WARNING: using validation set for testing!!!') test_data = load_dataset(cf, logger, test_ix, pp_data_path=cf.pp_test_data_path, pp_name=pp_name) logger.info("data set loaded with: {} test patients from {}".format(len(test_data.keys()), cf.pp_test_data_path)) batch_gen = {} batch_gen['test'] = PatientBatchIterator(test_data, cf=cf) batch_gen['n_test'] = len(test_data.keys()) if cf.max_test_patients=="all" else \ min(cf.max_test_patients, len(test_data.keys())) return batch_gen def load_dataset(cf, logger, subset_ixs=None, pp_data_path=None, pp_name=None): """ loads the dataset. if deployed in cloud also copies and unpacks the data to the working directory. :param subset_ixs: subset indices to be loaded from the dataset. used e.g. for testing to only load the test folds. :return: data: dictionary with one entry per patient (in this case per patient-breast, since they are treated as individual images for training) each entry is a dictionary containing respective meta-info as well as paths to the preprocessed numpy arrays to be loaded during batch-generation """ if pp_data_path is None: pp_data_path = cf.pp_data_path if pp_name is None: pp_name = cf.pp_name if cf.server_env: copy_data = True target_dir = os.path.join('/ssd', cf.slurm_job_id, pp_name) if not os.path.exists(target_dir): cf.data_source_dir = pp_data_path os.makedirs(target_dir) subprocess.call('rsync -av {} {}'.format( os.path.join(cf.data_source_dir, cf.input_df_name), os.path.join(target_dir, cf.input_df_name)), shell=True) logger.info('created target dir and info df at {}'.format(os.path.join(target_dir, cf.input_df_name))) elif subset_ixs is None: copy_data = False pp_data_path = target_dir p_df = pd.read_pickle(os.path.join(pp_data_path, cf.input_df_name)) if subset_ixs is not None: subset_pids = [np.unique(p_df.pid.tolist())[ix] for ix in subset_ixs] p_df = p_df[p_df.pid.isin(subset_pids)] logger.info('subset: selected {} instances from df'.format(len(p_df))) if cf.server_env: if copy_data: copy_and_unpack_data(logger, p_df.pid.tolist(), cf.fold_dir, cf.data_source_dir, target_dir) class_targets = p_df['class_id'].tolist() pids = p_df.pid.tolist() imgs = [os.path.join(pp_data_path, '{}.npy'.format(pid)) for pid in pids] segs = [os.path.join(pp_data_path,'{}.npy'.format(pid)) for pid in pids] data = OrderedDict() for ix, pid in enumerate(pids): data[pid] = {'data': imgs[ix], 'seg': segs[ix], 'pid': pid, 'class_target': [class_targets[ix]]} return data def create_data_gen_pipeline(patient_data, cf, do_aug=True): """ create mutli-threaded train/val/test batch generation and augmentation pipeline. :param patient_data: dictionary containing one dictionary per patient in the train/test subset. :param is_training: (optional) whether to perform data augmentation (training) or not (validation/testing) :return: multithreaded_generator """ # create instance of batch generator as first element in pipeline. data_gen = BatchGenerator(patient_data, batch_size=cf.batch_size, cf=cf) # add transformations to pipeline. my_transforms = [] if do_aug: mirror_transform = Mirror(axes=np.arange(2, cf.dim+2, 1)) my_transforms.append(mirror_transform) spatial_transform = SpatialTransform(patch_size=cf.patch_size[:cf.dim], patch_center_dist_from_border=cf.da_kwargs['rand_crop_dist'], do_elastic_deform=cf.da_kwargs['do_elastic_deform'], alpha=cf.da_kwargs['alpha'], sigma=cf.da_kwargs['sigma'], do_rotation=cf.da_kwargs['do_rotation'], angle_x=cf.da_kwargs['angle_x'], angle_y=cf.da_kwargs['angle_y'], angle_z=cf.da_kwargs['angle_z'], do_scale=cf.da_kwargs['do_scale'], scale=cf.da_kwargs['scale'], random_crop=cf.da_kwargs['random_crop']) my_transforms.append(spatial_transform) else: my_transforms.append(CenterCropTransform(crop_size=cf.patch_size[:cf.dim])) my_transforms.append(ConvertSegToBoundingBoxCoordinates(cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=cf.class_specific_seg_flag)) all_transforms = Compose(my_transforms) # multithreaded_generator = SingleThreadedAugmenter(data_gen, all_transforms) multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes=cf.n_workers, seeds=range(cf.n_workers)) return multithreaded_generator class BatchGenerator(SlimDataLoaderBase): """ creates the training/validation batch generator. Samples n_batch_size patients (draws a slice from each patient if 2D) from the data set while maintaining foreground-class balance. Returned patches are cropped/padded to pre_crop_size. Actual patch_size is obtained after data augmentation. :param data: data dictionary as provided by 'load_dataset'. :param batch_size: number of patients to sample for the batch :return dictionary containing the batch data (b, c, x, y, (z)) / seg (b, 1, x, y, (z)) / pids / class_target """ def __init__(self, data, batch_size, cf): super(BatchGenerator, self).__init__(data, batch_size) self.cf = cf def generate_train_batch(self): batch_data, batch_segs, batch_pids, batch_targets = [], [], [], [] class_targets_list = [v['class_target'] for (k, v) in self._data.items()] #samples patients towards equilibrium of foreground classes on a roi-level (after randomly sampling the ratio "batch_sample_slack). batch_ixs = dutils.get_class_balanced_patients( class_targets_list, self.batch_size, self.cf.head_classes - 1, slack_factor=self.cf.batch_sample_slack) patients = list(self._data.items()) for b in batch_ixs: patient = patients[b][1] all_data = np.load(patient['data'], mmap_mode='r') data = all_data[0] seg = all_data[1].astype('uint8') batch_pids.append(patient['pid']) batch_targets.append(patient['class_target']) batch_data.append(data[np.newaxis]) batch_segs.append(seg[np.newaxis]) data = np.array(batch_data) seg = np.array(batch_segs).astype(np.uint8) class_target = np.array(batch_targets) return {'data': data, 'seg': seg, 'pid': batch_pids, 'class_target': class_target} class PatientBatchIterator(SlimDataLoaderBase): """ creates a test generator that iterates over entire given dataset returning 1 patient per batch. Can be used for monitoring if cf.val_mode = 'patient_val' for a monitoring closer to actualy evaluation (done in 3D), if willing to accept speed-loss during training. :return: out_batch: dictionary containing one patient with batch_size = n_3D_patches in 3D or batch_size = n_2D_patches in 2D . """ def __init__(self, data, cf): #threads in augmenter super(PatientBatchIterator, self).__init__(data, 0) self.cf = cf self.patient_ix = 0 self.dataset_pids = [v['pid'] for (k, v) in data.items()] self.patch_size = cf.patch_size if len(self.patch_size) == 2: self.patch_size = self.patch_size + [1] def generate_train_batch(self): pid = self.dataset_pids[self.patient_ix] patient = self._data[pid] all_data = np.load(patient['data'], mmap_mode='r') data = all_data[0] seg = all_data[1].astype('uint8') batch_class_targets = np.array([patient['class_target']]) out_data = data[None, None] out_seg = seg[None, None] print('check patient data loader', out_data.shape, out_seg.shape) batch_2D = {'data': out_data, 'seg': out_seg, 'class_target': batch_class_targets, 'pid': pid} converter = ConvertSegToBoundingBoxCoordinates(dim=2, get_rois_from_seg_flag=False, class_specific_seg_flag=self.cf.class_specific_seg_flag) batch_2D = converter(**batch_2D) batch_2D.update({'patient_bb_target': batch_2D['bb_target'], 'patient_roi_labels': batch_2D['roi_labels'], 'original_img_shape': out_data.shape}) self.patient_ix += 1 if self.patient_ix == len(self.dataset_pids): self.patient_ix = 0 return batch_2D def copy_and_unpack_data(logger, pids, fold_dir, source_dir, target_dir): start_time = time.time() with open(os.path.join(fold_dir, 'file_list.txt'), 'w') as handle: for pid in pids: handle.write('{}.npy\n'.format(pid)) subprocess.call('rsync -av --files-from {} {} {}'.format(os.path.join(fold_dir, 'file_list.txt'), source_dir, target_dir), shell=True) # dutils.unpack_dataset(target_dir) copied_files = os.listdir(target_dir) logger.info("copying and unpacking data set finsihed : {} files in target dir: {}. took {} sec".format( len(copied_files), target_dir, np.round(time.time() - start_time, 0))) if __name__=="__main__": import utils.exp_utils as utils - from .configs import Configs + from configs import configs total_stime = time.time() - cf = Configs() - logger = utils.get_logger(0) + cf = configs() + logger = utils.get_logger("dev") batch_gen = get_train_generators(cf, logger) train_batch = next(batch_gen["train"]) mins, secs = divmod((time.time() - total_stime), 60) h, mins = divmod(mins, 60) t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs)) print("{} total runtime: {}".format(os.path.split(__file__)[1], t)) \ No newline at end of file diff --git a/models/mrcnn.py b/models/mrcnn.py index 0ba929c..55b1db0 100644 --- a/models/mrcnn.py +++ b/models/mrcnn.py @@ -1,1177 +1,1181 @@ #!/usr/bin/env python # Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """ Parts are based on https://github.com/multimodallearning/pytorch-mask-rcnn published under MIT license. """ import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.utils sys.path.append("..") import utils.model_utils as mutils import utils.exp_utils as utils from custom_extensions.nms import nms from custom_extensions.roi_align import roi_align ############################################################ # Networks on top of backbone ############################################################ class RPN(nn.Module): """ Region Proposal Network. """ def __init__(self, cf, conv): super(RPN, self).__init__() self.dim = conv.dim self.conv_shared = conv(cf.end_filts, cf.n_rpn_features, ks=3, stride=cf.rpn_anchor_stride, pad=1, relu=cf.relu) self.conv_class = conv(cf.n_rpn_features, 2 * len(cf.rpn_anchor_ratios), ks=1, stride=1, relu=None) self.conv_bbox = conv(cf.n_rpn_features, 2 * self.dim * len(cf.rpn_anchor_ratios), ks=1, stride=1, relu=None) def forward(self, x): """ :param x: input feature maps (b, in_channels, y, x, (z)) :return: rpn_class_logits (b, 2, n_anchors) :return: rpn_probs_logits (b, 2, n_anchors) :return: rpn_bbox (b, 2 * dim, n_anchors) """ # Shared convolutional base of the RPN. x = self.conv_shared(x) # Anchor Score. (batch, anchors per location * 2, y, x, (z)). rpn_class_logits = self.conv_class(x) # Reshape to (batch, 2, anchors) axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1) rpn_class_logits = rpn_class_logits.permute(*axes) rpn_class_logits = rpn_class_logits.contiguous() rpn_class_logits = rpn_class_logits.view(x.size()[0], -1, 2) # Softmax on last dimension (fg vs. bg). rpn_probs = F.softmax(rpn_class_logits, dim=2) # Bounding box refinement. (batch, anchors_per_location * (y, x, (z), log(h), log(w), (log(d)), y, x, (z)) rpn_bbox = self.conv_bbox(x) # Reshape to (batch, 2*dim, anchors) rpn_bbox = rpn_bbox.permute(*axes) rpn_bbox = rpn_bbox.contiguous() rpn_bbox = rpn_bbox.view(x.size()[0], -1, self.dim * 2) return [rpn_class_logits, rpn_probs, rpn_bbox] class Classifier(nn.Module): """ Head network for classification and bounding box refinement. Performs RoiAlign, processes resulting features through a shared convolutional base and finally branches off the classifier- and regression head. """ def __init__(self, cf, conv): super(Classifier, self).__init__() self.dim = conv.dim self.in_channels = cf.end_filts self.pool_size = cf.pool_size self.pyramid_levels = cf.pyramid_levels # instance_norm does not work with spatial dims (1, 1, (1)) norm = cf.norm if cf.norm != 'instance_norm' else None self.conv1 = conv(cf.end_filts, cf.end_filts * 4, ks=self.pool_size, stride=1, norm=norm, relu=cf.relu) self.conv2 = conv(cf.end_filts * 4, cf.end_filts * 4, ks=1, stride=1, norm=norm, relu=cf.relu) self.linear_class = nn.Linear(cf.end_filts * 4, cf.head_classes) self.linear_bbox = nn.Linear(cf.end_filts * 4, cf.head_classes * 2 * self.dim) def forward(self, x, rois): """ :param x: input feature maps (b, in_channels, y, x, (z)) :param rois: normalized box coordinates as proposed by the RPN to be forwarded through the second stage (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix). Proposals of all batch elements have been merged to one vector, while the origin info has been stored for re-allocation. :return: mrcnn_class_logits (n_proposals, n_head_classes) :return: mrcnn_bbox (n_proposals, n_head_classes, 2 * dim) predicted corrections to be applied to proposals for refinement. """ x = pyramid_roi_align(x, rois, self.pool_size, self.pyramid_levels, self.dim) x = self.conv1(x) x = self.conv2(x) x = x.view(-1, self.in_channels * 4) mrcnn_class_logits = self.linear_class(x) mrcnn_bbox = self.linear_bbox(x) mrcnn_bbox = mrcnn_bbox.view(mrcnn_bbox.size()[0], -1, self.dim * 2) return [mrcnn_class_logits, mrcnn_bbox] class Mask(nn.Module): """ Head network for proposal-based mask segmentation. Performs RoiAlign, some convolutions and applies sigmoid on the output logits to allow for overlapping classes. """ def __init__(self, cf, conv): super(Mask, self).__init__() self.pool_size = cf.mask_pool_size self.pyramid_levels = cf.pyramid_levels self.dim = conv.dim self.conv1 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu) self.conv2 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu) self.conv3 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu) self.conv4 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu) if conv.dim == 2: self.deconv = nn.ConvTranspose2d(cf.end_filts, cf.end_filts, kernel_size=2, stride=2) else: self.deconv = nn.ConvTranspose3d(cf.end_filts, cf.end_filts, kernel_size=2, stride=2) self.relu = nn.ReLU(inplace=True) if cf.relu == 'relu' else nn.LeakyReLU(inplace=True) self.conv5 = conv(cf.end_filts, cf.head_classes, ks=1, stride=1, relu=None) self.sigmoid = nn.Sigmoid() def forward(self, x, rois): """ :param x: input feature maps (b, in_channels, y, x, (z)) :param rois: normalized box coordinates as proposed by the RPN to be forwarded through the second stage (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix). Proposals of all batch elements have been merged to one vector, while the origin info has been stored for re-allocation. :return: x: masks (n_sampled_proposals (n_detections in inference), n_classes, y, x, (z)) """ x = pyramid_roi_align(x, rois, self.pool_size, self.pyramid_levels, self.dim) x = self.conv1(x) x = self.conv2(x) x = self.conv3(x) x = self.conv4(x) x = self.relu(self.deconv(x)) x = self.conv5(x) x = self.sigmoid(x) return x ############################################################ # Loss Functions ############################################################ def compute_rpn_class_loss(rpn_match, rpn_class_logits, shem_poolsize): """ :param rpn_match: (n_anchors). [-1, 0, 1] for negative, neutral, and positive matched anchors. :param rpn_class_logits: (n_anchors, 2). logits from RPN classifier. :param shem_poolsize: int. factor of top-k candidates to draw from per negative sample (stochastic-hard-example-mining). :return: loss: torch tensor :return: np_neg_ix: 1D array containing indices of the neg_roi_logits, which have been sampled for training. """ # filter out neutral anchors. pos_indices = torch.nonzero(rpn_match == 1) neg_indices = torch.nonzero(rpn_match == -1) # loss for positive samples if 0 not in pos_indices.size(): pos_indices = pos_indices.squeeze(1) roi_logits_pos = rpn_class_logits[pos_indices] pos_loss = F.cross_entropy(roi_logits_pos, torch.LongTensor([1] * pos_indices.shape[0]).cuda()) else: pos_loss = torch.FloatTensor([0]).cuda() # loss for negative samples: draw hard negative examples (SHEM) # that match the number of positive samples, but at least 1. if 0 not in neg_indices.size(): neg_indices = neg_indices.squeeze(1) roi_logits_neg = rpn_class_logits[neg_indices] negative_count = np.max((1, pos_indices.cpu().data.numpy().size)) roi_probs_neg = F.softmax(roi_logits_neg, dim=1) neg_ix = mutils.shem(roi_probs_neg, negative_count, shem_poolsize) neg_loss = F.cross_entropy(roi_logits_neg[neg_ix], torch.LongTensor([0] * neg_ix.shape[0]).cuda()) np_neg_ix = neg_ix.cpu().data.numpy() else: neg_loss = torch.FloatTensor([0]).cuda() np_neg_ix = np.array([]).astype('int32') loss = (pos_loss + neg_loss) / 2 return loss, np_neg_ix def compute_rpn_bbox_loss(rpn_target_deltas, rpn_pred_deltas, rpn_match): """ :param rpn_target_deltas: (b, n_positive_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd)))). Uses 0 padding to fill in unsed bbox deltas. :param rpn_pred_deltas: predicted deltas from RPN. (b, n_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd)))) :param rpn_match: (n_anchors). [-1, 0, 1] for negative, neutral, and positive matched anchors. :return: loss: torch 1D tensor. """ if 0 not in torch.nonzero(rpn_match == 1).size(): indices = torch.nonzero(rpn_match == 1).squeeze(1) # Pick bbox deltas that contribute to the loss rpn_pred_deltas = rpn_pred_deltas[indices] # Trim target bounding box deltas to the same length as rpn_bbox. target_deltas = rpn_target_deltas[:rpn_pred_deltas.size()[0], :] # Smooth L1 loss loss = F.smooth_l1_loss(rpn_pred_deltas, target_deltas) else: loss = torch.FloatTensor([0]).cuda() return loss def compute_mrcnn_class_loss(target_class_ids, pred_class_logits): """ :param target_class_ids: (n_sampled_rois) batch dimension was merged into roi dimension. :param pred_class_logits: (n_sampled_rois, n_classes) :return: loss: torch 1D tensor. """ if 0 not in target_class_ids.size(): loss = F.cross_entropy(pred_class_logits, target_class_ids.long()) else: loss = torch.FloatTensor([0.]).cuda() return loss def compute_mrcnn_bbox_loss(mrcnn_target_deltas, mrcnn_pred_deltas, target_class_ids): """ :param mrcnn_target_deltas: (n_sampled_rois, (dy, dx, (dz), log(dh), log(dw), (log(dh))) :param mrcnn_pred_deltas: (n_sampled_rois, n_classes, (dy, dx, (dz), log(dh), log(dw), (log(dh))) :param target_class_ids: (n_sampled_rois) :return: loss: torch 1D tensor. """ if 0 not in torch.nonzero(target_class_ids > 0).size(): positive_roi_ix = torch.nonzero(target_class_ids > 0)[:, 0] positive_roi_class_ids = target_class_ids[positive_roi_ix].long() target_bbox = mrcnn_target_deltas[positive_roi_ix, :].detach() pred_bbox = mrcnn_pred_deltas[positive_roi_ix, positive_roi_class_ids, :] loss = F.smooth_l1_loss(pred_bbox, target_bbox) else: loss = torch.FloatTensor([0]).cuda() return loss def compute_mrcnn_mask_loss(target_masks, pred_masks, target_class_ids): """ :param target_masks: (n_sampled_rois, y, x, (z)) A float32 tensor of values 0 or 1. Uses zero padding to fill array. :param pred_masks: (n_sampled_rois, n_classes, y, x, (z)) float32 tensor with values between [0, 1]. :param target_class_ids: (n_sampled_rois) :return: loss: torch 1D tensor. """ if 0 not in torch.nonzero(target_class_ids > 0).size(): # Only positive ROIs contribute to the loss. And only # the class specific mask of each ROI. positive_ix = torch.nonzero(target_class_ids > 0)[:, 0] positive_class_ids = target_class_ids[positive_ix].long() y_true = target_masks[positive_ix, :, :].detach() y_pred = pred_masks[positive_ix, positive_class_ids, :, :] loss = F.binary_cross_entropy(y_pred, y_true) else: loss = torch.FloatTensor([0]).cuda() return loss ############################################################ # Helper Layers ############################################################ def refine_proposals(rpn_pred_probs, rpn_pred_deltas, proposal_count, batch_anchors, cf): """ Receives anchor scores and selects a subset to pass as proposals to the second stage. Filtering is done based on anchor scores and non-max suppression to remove overlaps. It also applies bounding box refinment details to anchors. :param rpn_pred_probs: (b, n_anchors, 2) :param rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d)))) :return: batch_normalized_props: Proposals in normalized coordinates (b, proposal_count, (y1, x1, y2, x2, (z1), (z2), score)) :return: batch_out_proposals: Box coords + RPN foreground scores for monitoring/plotting (b, proposal_count, (y1, x1, y2, x2, (z1), (z2), score)) """ std_dev = torch.from_numpy(cf.rpn_bbox_std_dev[None]).float().cuda() norm = torch.from_numpy(cf.scale).float().cuda() anchors = batch_anchors.clone() batch_scores = rpn_pred_probs[:, :, 1] # norm deltas batch_deltas = rpn_pred_deltas * std_dev batch_normalized_props = [] batch_out_proposals = [] # loop over batch dimension. for ix in range(batch_scores.shape[0]): scores = batch_scores[ix] deltas = batch_deltas[ix] # improve performance by trimming to top anchors by score # and doing the rest on the smaller subset. pre_nms_limit = min(cf.pre_nms_limit, anchors.size()[0]) scores, order = scores.sort(descending=True) order = order[:pre_nms_limit] scores = scores[:pre_nms_limit] deltas = deltas[order, :] # apply deltas to anchors to get refined anchors and filter with non-maximum suppression. if batch_deltas.shape[-1] == 4: boxes = mutils.apply_box_deltas_2D(anchors[order, :], deltas) boxes = mutils.clip_boxes_2D(boxes, cf.window) else: boxes = mutils.apply_box_deltas_3D(anchors[order, :], deltas) boxes = mutils.clip_boxes_3D(boxes, cf.window) # boxes are y1,x1,y2,x2, torchvision-nms requires x1,y1,x2,y2, but consistent swap x<->y is irrelevant. keep = nms.nms(boxes, scores, cf.rpn_nms_threshold) keep = keep[:proposal_count] boxes = boxes[keep, :] rpn_scores = scores[keep][:, None] # pad missing boxes with 0. if boxes.shape[0] < proposal_count: n_pad_boxes = proposal_count - boxes.shape[0] zeros = torch.zeros([n_pad_boxes, boxes.shape[1]]).cuda() boxes = torch.cat([boxes, zeros], dim=0) zeros = torch.zeros([n_pad_boxes, rpn_scores.shape[1]]).cuda() rpn_scores = torch.cat([rpn_scores, zeros], dim=0) # concat box and score info for monitoring/plotting. batch_out_proposals.append(torch.cat((boxes, rpn_scores), 1).cpu().data.numpy()) # normalize dimensions to range of 0 to 1. normalized_boxes = boxes / norm assert torch.all(normalized_boxes <= 1), "normalized box coords >1 found" # add again batch dimension batch_normalized_props.append(normalized_boxes.unsqueeze(0)) batch_normalized_props = torch.cat(batch_normalized_props) batch_out_proposals = np.array(batch_out_proposals) return batch_normalized_props, batch_out_proposals def pyramid_roi_align(feature_maps, rois, pool_size, pyramid_levels, dim): """ Implements ROI Pooling on multiple levels of the feature pyramid. :param feature_maps: list of feature maps, each of shape (b, c, y, x , (z)) :param rois: proposals (normalized coords.) as returned by RPN. contain info about original batch element allocation. (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ixs) :param pool_size: list of poolsizes in dims: [x, y, (z)] :param pyramid_levels: list. [0, 1, 2, ...] :return: pooled: pooled feature map rois (n_proposals, c, poolsize_y, poolsize_x, (poolsize_z)) Output: Pooled regions in the shape: [num_boxes, height, width, channels]. The width and height are those specific in the pool_shape in the layer constructor. """ boxes = rois[:, :dim*2] batch_ixs = rois[:, dim*2] # Assign each ROI to a level in the pyramid based on the ROI area. if dim == 2: y1, x1, y2, x2 = boxes.chunk(4, dim=1) else: y1, x1, y2, x2, z1, z2 = boxes.chunk(6, dim=1) h = y2 - y1 w = x2 - x1 # Equation 1 in https://arxiv.org/abs/1612.03144. Account for # the fact that our coordinates are normalized here. # divide sqrt(h*w) by 1 instead image_area. roi_level = (4 + torch.log2(torch.sqrt(h*w))).round().int().clamp(pyramid_levels[0], pyramid_levels[-1]) # if Pyramid contains additional level P6, adapt the roi_level assignment accordingly. if len(pyramid_levels) == 5: roi_level[h*w > 0.65] = 5 # Loop through levels and apply ROI pooling to each. pooled = [] box_to_level = [] fmap_shapes = [f.shape for f in feature_maps] for level_ix, level in enumerate(pyramid_levels): ix = roi_level == level if not ix.any(): continue ix = torch.nonzero(ix)[:, 0] level_boxes = boxes[ix, :] # re-assign rois to feature map of original batch element. ind = batch_ixs[ix].int() # Keep track of which box is mapped to which level box_to_level.append(ix) # Stop gradient propogation to ROI proposals level_boxes = level_boxes.detach() if len(pool_size) == 2: # remap to feature map coordinate system y_exp, x_exp = fmap_shapes[level_ix][2:] # exp = expansion level_boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp], dtype=torch.float32).cuda()) pooled_features = roi_align.roi_align_2d(feature_maps[level_ix], torch.cat((ind.unsqueeze(1).float(), level_boxes), dim=1), pool_size) else: y_exp, x_exp, z_exp = fmap_shapes[level_ix][2:] level_boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp, z_exp, z_exp], dtype=torch.float32).cuda()) pooled_features = roi_align.roi_align_3d(feature_maps[level_ix], torch.cat((ind.unsqueeze(1).float(), level_boxes), dim=1), pool_size) pooled.append(pooled_features) # Pack pooled features into one tensor pooled = torch.cat(pooled, dim=0) # Pack box_to_level mapping into one array and add another # column representing the order of pooled boxes box_to_level = torch.cat(box_to_level, dim=0) # Rearrange pooled features to match the order of the original boxes _, box_to_level = torch.sort(box_to_level) pooled = pooled[box_to_level, :, :] return pooled def detection_target_layer(batch_proposals, batch_mrcnn_class_scores, batch_gt_class_ids, batch_gt_boxes, batch_gt_masks, cf): """ Subsamples proposals for mrcnn losses and generates targets. Sampling is done per batch element, seems to have positive effects on training, as opposed to sampling over entire batch. Negatives are sampled via stochastic-hard-example-mining (SHEM), where a number of negative proposals are drawn from larger pool of highest scoring proposals for stochasticity. Scoring is obtained here as the max over all foreground probabilities as returned by mrcnn_classifier (worked better than loss-based class balancing methods like "online-hard-example-mining" or "focal loss".) :param batch_proposals: (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ixs). boxes as proposed by RPN. n_proposals here is determined by batch_size * POST_NMS_ROIS. :param batch_mrcnn_class_scores: (n_proposals, n_classes) :param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels. :param batch_gt_boxes: list over batch elements. Each element is a list over the corresponding roi target coordinates. :param batch_gt_masks: list over batch elements. Each element is binary mask of shape (n_gt_rois, y, x, (z), c) :return: sample_indices: (n_sampled_rois) indices of sampled proposals to be used for loss functions. :return: target_class_ids: (n_sampled_rois)containing target class labels of sampled proposals. :return: target_deltas: (n_sampled_rois, 2 * dim) containing target deltas of sampled proposals for box refinement. :return: target_masks: (n_sampled_rois, y, x, (z)) containing target masks of sampled proposals. """ # normalization of target coordinates if cf.dim == 2: h, w = cf.patch_size scale = torch.from_numpy(np.array([h, w, h, w])).float().cuda() else: h, w, z = cf.patch_size scale = torch.from_numpy(np.array([h, w, h, w, z, z])).float().cuda() positive_count = 0 negative_count = 0 sample_positive_indices = [] sample_negative_indices = [] sample_deltas = [] sample_masks = [] sample_class_ids = [] std_dev = torch.from_numpy(cf.bbox_std_dev).float().cuda() # loop over batch and get positive and negative sample rois. for b in range(len(batch_gt_class_ids)): gt_class_ids = torch.from_numpy(batch_gt_class_ids[b]).int().cuda() gt_masks = torch.from_numpy(batch_gt_masks[b]).float().cuda() if np.any(batch_gt_class_ids[b] > 0): # skip roi selection for no gt images. gt_boxes = torch.from_numpy(batch_gt_boxes[b]).float().cuda() / scale else: gt_boxes = torch.FloatTensor().cuda() # get proposals and indices of current batch element. proposals = batch_proposals[batch_proposals[:, -1] == b][:, :-1] batch_element_indices = torch.nonzero(batch_proposals[:, -1] == b).squeeze(1) # Compute overlaps matrix [proposals, gt_boxes] if 0 not in gt_boxes.size(): if gt_boxes.shape[1] == 4: assert cf.dim == 2, "gt_boxes shape {} doesnt match cf.dim{}".format(gt_boxes.shape, cf.dim) overlaps = mutils.bbox_overlaps_2D(proposals, gt_boxes) else: assert cf.dim == 3, "gt_boxes shape {} doesnt match cf.dim{}".format(gt_boxes.shape, cf.dim) overlaps = mutils.bbox_overlaps_3D(proposals, gt_boxes) # Determine postive and negative ROIs roi_iou_max = torch.max(overlaps, dim=1)[0] # 1. Positive ROIs are those with >= 0.5 IoU with a GT box positive_roi_bool = roi_iou_max >= (0.5 if cf.dim == 2 else 0.3) # 2. Negative ROIs are those with < 0.1 with every GT box. negative_roi_bool = roi_iou_max < (0.1 if cf.dim == 2 else 0.01) else: positive_roi_bool = torch.FloatTensor().cuda() negative_roi_bool = torch.from_numpy(np.array([1]*proposals.shape[0])).cuda() # Sample Positive ROIs if 0 not in torch.nonzero(positive_roi_bool).size(): positive_indices = torch.nonzero(positive_roi_bool).squeeze(1) positive_samples = int(cf.train_rois_per_image * cf.roi_positive_ratio) rand_idx = torch.randperm(positive_indices.size()[0]) rand_idx = rand_idx[:positive_samples].cuda() positive_indices = positive_indices[rand_idx] positive_samples = positive_indices.size()[0] positive_rois = proposals[positive_indices, :] # Assign positive ROIs to GT boxes. positive_overlaps = overlaps[positive_indices, :] roi_gt_box_assignment = torch.max(positive_overlaps, dim=1)[1] roi_gt_boxes = gt_boxes[roi_gt_box_assignment, :] roi_gt_class_ids = gt_class_ids[roi_gt_box_assignment] # Compute bbox refinement targets for positive ROIs deltas = mutils.box_refinement(positive_rois, roi_gt_boxes) deltas /= std_dev # Assign positive ROIs to GT masks - roi_masks = gt_masks[roi_gt_box_assignment].unsqueeze(1) - assert roi_masks.shape[-1] == 1 + roi_masks = gt_masks[roi_gt_box_assignment] + assert roi_masks.shape[1] == 1, "desired to have more than one channel in gt masks?" # Compute mask targets boxes = positive_rois box_ids = torch.arange(roi_masks.shape[0]).cuda().unsqueeze(1).float() - if len(cf.mask_shape) == 2: - # todo what are the dims of roi_masks? (n_matched_boxes_with_gts, 1 (dummy channel dim), y,x, 1 (WHY?)) + # need to remap normalized box coordinates to unnormalized mask coordinates. + y_exp, x_exp = roi_masks.shape[2:] # exp = expansion + boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp], dtype=torch.float32).cuda()) masks = roi_align.roi_align_2d(roi_masks, torch.cat((box_ids, boxes), dim=1), cf.mask_shape) else: + y_exp, x_exp, z_exp = roi_masks.shape[2:] # exp = expansion + boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp, z_exp, z_exp], dtype=torch.float32).cuda()) masks = roi_align.roi_align_3d(roi_masks, torch.cat((box_ids, boxes), dim=1), cf.mask_shape) - masks = masks.squeeze(1) # Threshold mask pixels at 0.5 to have GT masks be 0 or 1 to use with # binary cross entropy loss. masks = torch.round(masks) sample_positive_indices.append(batch_element_indices[positive_indices]) sample_deltas.append(deltas) sample_masks.append(masks) sample_class_ids.append(roi_gt_class_ids) positive_count += positive_samples else: positive_samples = 0 # Negative ROIs. Add enough to maintain positive:negative ratio, but at least 1. Sample via SHEM. if 0 not in torch.nonzero(negative_roi_bool).size(): negative_indices = torch.nonzero(negative_roi_bool).squeeze(1) r = 1.0 / cf.roi_positive_ratio b_neg_count = np.max((int(r * positive_samples - positive_samples), 1)) roi_probs_neg = batch_mrcnn_class_scores[batch_element_indices[negative_indices]] raw_sampled_indices = mutils.shem(roi_probs_neg, b_neg_count, cf.shem_poolsize) sample_negative_indices.append(batch_element_indices[negative_indices[raw_sampled_indices]]) negative_count += raw_sampled_indices.size()[0] if len(sample_positive_indices) > 0: target_deltas = torch.cat(sample_deltas) target_masks = torch.cat(sample_masks) target_class_ids = torch.cat(sample_class_ids) # Pad target information with zeros for negative ROIs. if positive_count > 0 and negative_count > 0: sample_indices = torch.cat((torch.cat(sample_positive_indices), torch.cat(sample_negative_indices)), dim=0) zeros = torch.zeros(negative_count).int().cuda() target_class_ids = torch.cat([target_class_ids, zeros], dim=0) zeros = torch.zeros(negative_count, cf.dim * 2).cuda() target_deltas = torch.cat([target_deltas, zeros], dim=0) zeros = torch.zeros(negative_count, *cf.mask_shape).cuda() target_masks = torch.cat([target_masks, zeros], dim=0) elif positive_count > 0: sample_indices = torch.cat(sample_positive_indices) elif negative_count > 0: sample_indices = torch.cat(sample_negative_indices) zeros = torch.zeros(negative_count).int().cuda() target_class_ids = zeros zeros = torch.zeros(negative_count, cf.dim * 2).cuda() target_deltas = zeros zeros = torch.zeros(negative_count, *cf.mask_shape).cuda() target_masks = zeros else: sample_indices = torch.LongTensor().cuda() target_class_ids = torch.IntTensor().cuda() target_deltas = torch.FloatTensor().cuda() target_masks = torch.FloatTensor().cuda() return sample_indices, target_class_ids, target_deltas, target_masks ############################################################ # Output Handler ############################################################ # def refine_detections(rois, probs, deltas, batch_ixs, cf): # """ # Refine classified proposals, filter overlaps and return final detections. # # :param rois: (n_proposals, 2 * dim) normalized boxes as proposed by RPN. n_proposals = batch_size * POST_NMS_ROIS # :param probs: (n_proposals, n_classes) softmax probabilities for all rois as predicted by mrcnn classifier. # :param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by mrcnn bbox regressor. # :param batch_ixs: (n_proposals) batch element assignemnt info for re-allocation. # :return: result: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)) # """ # # class IDs per ROI. Since scores of all classes are of interest (not just max class), all are kept at this point. # class_ids = [] # fg_classes = cf.head_classes - 1 # # repeat vectors to fill in predictions for all foreground classes. # for ii in range(1, fg_classes + 1): # class_ids += [ii] * rois.shape[0] # class_ids = torch.from_numpy(np.array(class_ids)).cuda() # # rois = rois.repeat(fg_classes, 1) # probs = probs.repeat(fg_classes, 1) # deltas = deltas.repeat(fg_classes, 1, 1) # batch_ixs = batch_ixs.repeat(fg_classes) # # # get class-specific scores and bounding box deltas # idx = torch.arange(class_ids.size()[0]).long().cuda() # class_scores = probs[idx, class_ids] # deltas_specific = deltas[idx, class_ids] # batch_ixs = batch_ixs[idx] # # # apply bounding box deltas. re-scale to image coordinates. # std_dev = torch.from_numpy(np.reshape(cf.rpn_bbox_std_dev, [1, cf.dim * 2])).float().cuda() # scale = torch.from_numpy(cf.scale).float().cuda() # refined_rois = mutils.apply_box_deltas_2D(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \ # mutils.apply_box_deltas_3D(rois, deltas_specific * std_dev) * scale # # # round and cast to int since we're deadling with pixels now # refined_rois = mutils.clip_to_window(cf.window, refined_rois) # refined_rois = torch.round(refined_rois) # # # filter out low confidence boxes # keep = idx # keep_bool = (class_scores >= cf.model_min_confidence) # if 0 not in torch.nonzero(keep_bool).size(): # # score_keep = torch.nonzero(keep_bool)[:, 0] # pre_nms_class_ids = class_ids[score_keep] # pre_nms_rois = refined_rois[score_keep] # pre_nms_scores = class_scores[score_keep] # pre_nms_batch_ixs = batch_ixs[score_keep] # # for j, b in enumerate(mutils.unique1d(pre_nms_batch_ixs)): # # bixs = torch.nonzero(pre_nms_batch_ixs == b)[:, 0] # bix_class_ids = pre_nms_class_ids[bixs] # bix_rois = pre_nms_rois[bixs] # bix_scores = pre_nms_scores[bixs] # # for i, class_id in enumerate(mutils.unique1d(bix_class_ids)): # # ixs = torch.nonzero(bix_class_ids == class_id)[:, 0] # # nms expects boxes sorted by score. # ix_rois = bix_rois[ixs] # ix_scores = bix_scores[ixs] # ix_scores, order = ix_scores.sort(descending=True) # ix_rois = ix_rois[order, :] # # if cf.dim == 2: # class_keep = nms_2D(torch.cat((ix_rois, ix_scores.unsqueeze(1)), dim=1), cf.detection_nms_threshold) # else: # class_keep = nms_3D(torch.cat((ix_rois, ix_scores.unsqueeze(1)), dim=1), cf.detection_nms_threshold) # # # map indices back. # class_keep = keep[score_keep[bixs[ixs[order[class_keep]]]]] # # merge indices over classes for current batch element # b_keep = class_keep if i == 0 else mutils.unique1d(torch.cat((b_keep, class_keep))) # # # only keep top-k boxes of current batch-element # top_ids = class_scores[b_keep].sort(descending=True)[1][:cf.model_max_instances_per_batch_element] # b_keep = b_keep[top_ids] # # # merge indices over batch elements. # batch_keep = b_keep if j == 0 else mutils.unique1d(torch.cat((batch_keep, b_keep))) # # keep = batch_keep # # else: # keep = torch.tensor([0]).long().cuda() # # # arrange output # result = torch.cat((refined_rois[keep], # batch_ixs[keep].unsqueeze(1), # class_ids[keep].unsqueeze(1).float(), # class_scores[keep].unsqueeze(1)), dim=1) # # return result def refine_detections(cf, batch_ixs, rois, deltas, scores): """ Refine classified proposals (apply deltas to rpn rois), filter overlaps (nms) and return final detections. :param rois: (n_proposals, 2 * dim) normalized boxes as proposed by RPN. n_proposals = batch_size * POST_NMS_ROIS :param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by mrcnn bbox regressor. :param batch_ixs: (n_proposals) batch element assignment info for re-allocation. :param scores: (n_proposals, n_classes) probabilities for all classes per roi as predicted by mrcnn classifier. :return: result: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score, *regression vector features)) """ # class IDs per ROI. Since scores of all classes are of interest (not just max class), all are kept at this point. class_ids = [] fg_classes = cf.head_classes - 1 # repeat vectors to fill in predictions for all foreground classes. for ii in range(1, fg_classes + 1): class_ids += [ii] * rois.shape[0] class_ids = torch.from_numpy(np.array(class_ids)).cuda() batch_ixs = batch_ixs.repeat(fg_classes) rois = rois.repeat(fg_classes, 1) deltas = deltas.repeat(fg_classes, 1, 1) scores = scores.repeat(fg_classes, 1) # get class-specific scores and bounding box deltas idx = torch.arange(class_ids.size()[0]).long().cuda() # using idx instead of slice [:,] squashes first dimension. #len(class_ids)>scores.shape[1] --> probs is broadcasted by expansion from fg_classes-->len(class_ids) batch_ixs = batch_ixs[idx] deltas_specific = deltas[idx, class_ids] class_scores = scores[idx, class_ids] # apply bounding box deltas. re-scale to image coordinates. std_dev = torch.from_numpy(np.reshape(cf.rpn_bbox_std_dev, [1, cf.dim * 2])).float().cuda() scale = torch.from_numpy(cf.scale).float().cuda() refined_rois = mutils.apply_box_deltas_2D(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \ mutils.apply_box_deltas_3D(rois, deltas_specific * std_dev) * scale # round and cast to int since we're dealing with pixels now refined_rois = mutils.clip_to_window(cf.window, refined_rois) refined_rois = torch.round(refined_rois) # filter out low confidence boxes keep = idx keep_bool = (class_scores >= cf.model_min_confidence) if not 0 in torch.nonzero(keep_bool).size(): score_keep = torch.nonzero(keep_bool)[:, 0] pre_nms_class_ids = class_ids[score_keep] pre_nms_rois = refined_rois[score_keep] pre_nms_scores = class_scores[score_keep] pre_nms_batch_ixs = batch_ixs[score_keep] for j, b in enumerate(mutils.unique1d(pre_nms_batch_ixs)): bixs = torch.nonzero(pre_nms_batch_ixs == b)[:, 0] bix_class_ids = pre_nms_class_ids[bixs] bix_rois = pre_nms_rois[bixs] bix_scores = pre_nms_scores[bixs] for i, class_id in enumerate(mutils.unique1d(bix_class_ids)): ixs = torch.nonzero(bix_class_ids == class_id)[:, 0] # nms expects boxes sorted by score. ix_rois = bix_rois[ixs] ix_scores = bix_scores[ixs] ix_scores, order = ix_scores.sort(descending=True) ix_rois = ix_rois[order, :] class_keep = nms.nms(ix_rois, ix_scores, cf.detection_nms_threshold) # map indices back. class_keep = keep[score_keep[bixs[ixs[order[class_keep]]]]] # merge indices over classes for current batch element b_keep = class_keep if i == 0 else mutils.unique1d(torch.cat((b_keep, class_keep))) # only keep top-k boxes of current batch-element top_ids = class_scores[b_keep].sort(descending=True)[1][:cf.model_max_instances_per_batch_element] b_keep = b_keep[top_ids] # merge indices over batch elements. batch_keep = b_keep if j == 0 else mutils.unique1d(torch.cat((batch_keep, b_keep))) keep = batch_keep else: keep = torch.tensor([0]).long().cuda() # arrange output output = [refined_rois[keep], batch_ixs[keep].unsqueeze(1)] output += [class_ids[keep].unsqueeze(1).float(), class_scores[keep].unsqueeze(1)] result = torch.cat(output, dim=1) # shape: (n_keeps, catted feats), catted feats: [0:dim*2] are box_coords, [dim*2] are batch_ics, # [dim*2+1] are class_ids, [dim*2+2] are scores, [dim*2+3:] are regression vector features (incl uncertainty) return result def get_results(cf, img_shape, detections, detection_masks, box_results_list=None, return_masks=True): """ Restores batch dimension of merged detections, unmolds detections, creates and fills results dict. :param img_shape: :param detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score) :param detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head. :param box_results_list: None or list of output boxes for monitoring/plotting. each element is a list of boxes per batch element. :param return_masks: boolean. If True, full resolution masks are returned for all proposals (speed trade-off). :return: results_dict: dictionary with keys: 'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary: [[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...] 'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, 1] only fg. vs. bg for now. class-specific return of masks will come with implementation of instance segmentation evaluation. """ detections = detections.cpu().data.numpy() if cf.dim == 2: detection_masks = detection_masks.permute(0, 2, 3, 1).cpu().data.numpy() else: detection_masks = detection_masks.permute(0, 2, 3, 4, 1).cpu().data.numpy() # restore batch dimension of merged detections using the batch_ix info. batch_ixs = detections[:, cf.dim*2] detections = [detections[batch_ixs == ix] for ix in range(img_shape[0])] mrcnn_mask = [detection_masks[batch_ixs == ix] for ix in range(img_shape[0])] # for test_forward, where no previous list exists. if box_results_list is None: box_results_list = [[] for _ in range(img_shape[0])] seg_preds = [] # loop over batch and unmold detections. for ix in range(img_shape[0]): if 0 not in detections[ix].shape: boxes = detections[ix][:, :2 * cf.dim].astype(np.int32) class_ids = detections[ix][:, 2 * cf.dim + 1].astype(np.int32) scores = detections[ix][:, 2 * cf.dim + 2] masks = mrcnn_mask[ix][np.arange(boxes.shape[0]), ..., class_ids] # Filter out detections with zero area. Often only happens in early # stages of training when the network weights are still a bit random. if cf.dim == 2: exclude_ix = np.where((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) <= 0)[0] else: exclude_ix = np.where( (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 5] - boxes[:, 4]) <= 0)[0] if exclude_ix.shape[0] > 0: boxes = np.delete(boxes, exclude_ix, axis=0) class_ids = np.delete(class_ids, exclude_ix, axis=0) scores = np.delete(scores, exclude_ix, axis=0) masks = np.delete(masks, exclude_ix, axis=0) # Resize masks to original image size and set boundary threshold. full_masks = [] permuted_image_shape = list(img_shape[2:]) + [img_shape[1]] if return_masks: for i in range(masks.shape[0]): # Convert neural network mask to full size mask. full_masks.append(mutils.unmold_mask_2D(masks[i], boxes[i], permuted_image_shape) if cf.dim == 2 else mutils.unmold_mask_3D(masks[i], boxes[i], permuted_image_shape)) # if masks are returned, take max over binary full masks of all predictions in this image. - # right now only binary masks for plotting/monitoring. for instance segmentation return all proposal maks. + # right now only binary masks for plotting/monitoring. for instance segmentation return all proposal masks. final_masks = np.max(np.array(full_masks), 0) if len(full_masks) > 0 else np.zeros( (*permuted_image_shape[:-1],)) - # add final perdictions to results. + # add final predictions to results. if 0 not in boxes.shape: for ix2, score in enumerate(scores): box_results_list[ix].append({'box_coords': boxes[ix2], 'box_score': score, 'box_type': 'det', 'box_pred_class_id': class_ids[ix2]}) else: # pad with zero dummy masks. final_masks = np.zeros(img_shape[2:]) seg_preds.append(final_masks) # create and fill results dictionary. results_dict = {'boxes': box_results_list, 'seg_preds': np.round(np.array(seg_preds))[:, np.newaxis].astype('uint8')} return results_dict ############################################################ # Mask R-CNN Class ############################################################ class net(nn.Module): def __init__(self, cf, logger): super(net, self).__init__() self.cf = cf self.logger = logger self.build() if self.cf.weight_init is not None: logger.info("using pytorch weight init of type {}".format(self.cf.weight_init)) mutils.initialize_weights(self) else: logger.info("using default pytorch weight init") def build(self): """Build Mask R-CNN architecture.""" # Image size must be dividable by 2 multiple times. h, w = self.cf.patch_size[:2] if h / 2**5 != int(h / 2**5) or w / 2**5 != int(w / 2**5): raise Exception("Image size must be dividable by 2 at least 5 times " "to avoid fractions when downscaling and upscaling." "For example, use 256, 320, 384, 448, 512, ... etc. ") if len(self.cf.patch_size) == 3: d = self.cf.patch_size[2] if d / 2**3 != int(d / 2**3): raise Exception("Image z dimension must be dividable by 2 at least 3 times " "to avoid fractions when downscaling and upscaling.") # instanciate abstract multi dimensional conv class and backbone class. conv = mutils.NDConvGenerator(self.cf.dim) backbone = utils.import_module('bbone', self.cf.backbone_path) # build Anchors, FPN, RPN, Classifier / Bbox-Regressor -head, Mask-head self.np_anchors = mutils.generate_pyramid_anchors(self.logger, self.cf) self.anchors = torch.from_numpy(self.np_anchors).float().cuda() self.fpn = backbone.FPN(self.cf, conv) self.rpn = RPN(self.cf, conv) self.classifier = Classifier(self.cf, conv) self.mask = Mask(self.cf, conv) def train_forward(self, batch, is_validation=False): """ train method (also used for validation monitoring). wrapper around forward pass of network. prepares input data for processing, computes losses, and stores outputs in a dictionary. :param batch: dictionary containing 'data', 'seg', etc. + data_dict['roi_masks']: (b, n(b), 1, h(n), w(n) (z(n))) list like batch['roi_labels'] but with + arrays (masks) inplace of integers. n == number of rois per this batch element. :return: results_dict: dictionary with keys: 'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary: [[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...] 'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, n_classes]. 'monitor_values': dict of values to be monitored. """ img = batch['data'] gt_class_ids = batch['roi_labels'] gt_boxes = batch['bb_target'] - axes = (0, 2, 3, 1) if self.cf.dim == 2 else (0, 2, 3, 4, 1) - gt_masks = [np.transpose(batch['roi_masks'][ii], axes=axes) for ii in range(len(batch['roi_masks']))] - - + #axes = (0, 2, 3, 1) if self.cf.dim == 2 else (0, 2, 3, 4, 1) + #gt_masks = [np.transpose(batch['roi_masks'][ii], axes=axes) for ii in range(len(batch['roi_masks']))] + # --> now GT masks has c==channels in last dimension. + gt_masks = batch['roi_masks'] img = torch.from_numpy(img).float().cuda() batch_rpn_class_loss = torch.FloatTensor([0]).cuda() batch_rpn_bbox_loss = torch.FloatTensor([0]).cuda() # list of output boxes for monitoring/plotting. each element is a list of boxes per batch element. box_results_list = [[] for _ in range(img.shape[0])] #forward passes. 1. general forward pass, where no activations are saved in second stage (for performance # monitoring and loss sampling). 2. second stage forward pass of sampled rois with stored activations for backprop. rpn_class_logits, rpn_pred_deltas, proposal_boxes, detections, detection_masks = self.forward(img) mrcnn_class_logits, mrcnn_pred_deltas, mrcnn_pred_mask, target_class_ids, mrcnn_target_deltas, target_mask, \ sample_proposals = self.loss_samples_forward(gt_class_ids, gt_boxes, gt_masks) # loop over batch for b in range(img.shape[0]): if len(gt_boxes[b]) > 0: # add gt boxes to output list for monitoring. for ix in range(len(gt_boxes[b])): box_results_list[b].append({'box_coords': batch['bb_target'][b][ix], 'box_label': batch['roi_labels'][b][ix], 'box_type': 'gt'}) # match gt boxes with anchors to generate targets for RPN losses. rpn_match, rpn_target_deltas = mutils.gt_anchor_matching(self.cf, self.np_anchors, gt_boxes[b]) # add positive anchors used for loss to output list for monitoring. pos_anchors = mutils.clip_boxes_numpy(self.np_anchors[np.argwhere(rpn_match == 1)][:, 0], img.shape[2:]) for p in pos_anchors: box_results_list[b].append({'box_coords': p, 'box_type': 'pos_anchor'}) else: rpn_match = np.array([-1]*self.np_anchors.shape[0]) rpn_target_deltas = np.array([0]) rpn_match_gpu = torch.from_numpy(rpn_match).cuda() rpn_target_deltas = torch.from_numpy(rpn_target_deltas).float().cuda() # compute RPN losses. rpn_class_loss, neg_anchor_ix = compute_rpn_class_loss(rpn_match_gpu, rpn_class_logits[b], self.cf.shem_poolsize) rpn_bbox_loss = compute_rpn_bbox_loss(rpn_target_deltas, rpn_pred_deltas[b], rpn_match_gpu) batch_rpn_class_loss += rpn_class_loss / img.shape[0] batch_rpn_bbox_loss += rpn_bbox_loss / img.shape[0] # add negative anchors used for loss to output list for monitoring. neg_anchors = mutils.clip_boxes_numpy(self.np_anchors[rpn_match == -1][neg_anchor_ix], img.shape[2:]) for n in neg_anchors: box_results_list[b].append({'box_coords': n, 'box_type': 'neg_anchor'}) # add highest scoring proposals to output list for monitoring. rpn_proposals = proposal_boxes[b][proposal_boxes[b, :, -1].argsort()][::-1] for r in rpn_proposals[:self.cf.n_plot_rpn_props, :-1]: box_results_list[b].append({'box_coords': r, 'box_type': 'prop'}) # add positive and negative roi samples used for mrcnn losses to output list for monitoring. if 0 not in sample_proposals.shape: rois = mutils.clip_to_window(self.cf.window, sample_proposals).cpu().data.numpy() for ix, r in enumerate(rois): box_results_list[int(r[-1])].append({'box_coords': r[:-1] * self.cf.scale, 'box_type': 'pos_class' if target_class_ids[ix] > 0 else 'neg_class'}) batch_rpn_class_loss = batch_rpn_class_loss batch_rpn_bbox_loss = batch_rpn_bbox_loss # compute mrcnn losses. mrcnn_class_loss = compute_mrcnn_class_loss(target_class_ids, mrcnn_class_logits) mrcnn_bbox_loss = compute_mrcnn_bbox_loss(mrcnn_target_deltas, mrcnn_pred_deltas, target_class_ids) # mrcnn can be run without pixelwise annotations available (Faster R-CNN mode). # In this case, the mask_loss is taken out of training. if not self.cf.frcnn_mode: mrcnn_mask_loss = compute_mrcnn_mask_loss(target_mask, mrcnn_pred_mask, target_class_ids) else: mrcnn_mask_loss = torch.FloatTensor([0]).cuda() loss = batch_rpn_class_loss + batch_rpn_bbox_loss + mrcnn_class_loss + mrcnn_bbox_loss + mrcnn_mask_loss # monitor RPN performance: detection count = the number of correctly matched proposals per fg-class. dcount = [list(target_class_ids.cpu().data.numpy()).count(c) for c in np.arange(self.cf.head_classes)[1:]] # run unmolding of predictions for monitoring and merge all results to one dictionary. return_masks = self.cf.return_masks_in_val if is_validation else False results_dict = get_results(self.cf, img.shape, detections, detection_masks, box_results_list, return_masks=return_masks) results_dict['torch_loss'] = loss results_dict['monitor_values'] = {'loss': loss.item(), 'class_loss': mrcnn_class_loss.item()} results_dict['logger_string'] = \ "loss: {0:.2f}, rpn_class: {1:.2f}, rpn_bbox: {2:.2f}, mrcnn_class: {3:.2f}, mrcnn_bbox: {4:.2f}, " \ "mrcnn_mask: {5:.2f}, dcount {6}".format(loss.item(), batch_rpn_class_loss.item(), batch_rpn_bbox_loss.item(), mrcnn_class_loss.item(), mrcnn_bbox_loss.item(), mrcnn_mask_loss.item(), dcount) return results_dict def test_forward(self, batch, return_masks=True): """ test method. wrapper around forward pass of network without usage of any ground truth information. prepares input data for processing and stores outputs in a dictionary. :param batch: dictionary containing 'data' :param return_masks: boolean. If True, full resolution masks are returned for all proposals (speed trade-off). :return: results_dict: dictionary with keys: 'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary: [[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...] 'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, n_classes] """ img = batch['data'] img = torch.from_numpy(img).float().cuda() _, _, _, detections, detection_masks = self.forward(img) results_dict = get_results(self.cf, img.shape, detections, detection_masks, return_masks=return_masks) return results_dict def forward(self, img, is_training=True): """ :param img: input images (b, c, y, x, (z)). :return: rpn_pred_logits: (b, n_anchors, 2) :return: rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d)))) :return: batch_proposal_boxes: (b, n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix)) only for monitoring/plotting. :return: detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score) :return: detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head. """ # extract features. fpn_outs = self.fpn(img) rpn_feature_maps = [fpn_outs[i] for i in self.cf.pyramid_levels] self.mrcnn_feature_maps = rpn_feature_maps # loop through pyramid layers and apply RPN. layer_outputs = [] # list of lists for p in rpn_feature_maps: layer_outputs.append(self.rpn(p)) # concatenate layer outputs. # convert from list of lists of level outputs to list of lists of outputs across levels. # e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]] outputs = list(zip(*layer_outputs)) outputs = [torch.cat(list(o), dim=1) for o in outputs] rpn_pred_logits, rpn_pred_probs, rpn_pred_deltas = outputs # generate proposals: apply predicted deltas to anchors and filter by foreground scores from RPN classifier. proposal_count = self.cf.post_nms_rois_training if is_training else self.cf.post_nms_rois_inference batch_rpn_rois, batch_proposal_boxes = refine_proposals(rpn_pred_probs, rpn_pred_deltas, proposal_count, self.anchors, self.cf) # merge batch dimension of proposals while storing allocation info in coordinate dimension. batch_ixs = torch.from_numpy(np.repeat(np.arange(batch_rpn_rois.shape[0]), batch_rpn_rois.shape[1])).float().cuda() rpn_rois = batch_rpn_rois.view(-1, batch_rpn_rois.shape[2]) self.rpn_rois_batch_info = torch.cat((rpn_rois, batch_ixs.unsqueeze(1)), dim=1) # this is the first of two forward passes in the second stage, where no activations are stored for backprop. # here, all proposals are forwarded (with virtual_batch_size = batch_size * post_nms_rois.) # for inference/monitoring as well as sampling of rois for the loss functions. # processed in chunks of roi_chunk_size to re-adjust to gpu-memory. chunked_rpn_rois = self.rpn_rois_batch_info.split(self.cf.roi_chunk_size) class_logits_list, bboxes_list = [], [] with torch.no_grad(): for chunk in chunked_rpn_rois: chunk_class_logits, chunk_bboxes = self.classifier(self.mrcnn_feature_maps, chunk) class_logits_list.append(chunk_class_logits) bboxes_list.append(chunk_bboxes) batch_mrcnn_class_logits = torch.cat(class_logits_list, 0) batch_mrcnn_bbox = torch.cat(bboxes_list, 0) self.batch_mrcnn_class_scores = F.softmax(batch_mrcnn_class_logits, dim=1) # refine classified proposals, filter and return final detections. detections = refine_detections(self.cf, batch_ixs, rpn_rois, batch_mrcnn_bbox, self.batch_mrcnn_class_scores) # forward remaining detections through mask-head to generate corresponding masks. scale = [img.shape[2]] * 4 + [img.shape[-1]] * 2 scale = torch.from_numpy(np.array(scale[:self.cf.dim * 2] + [1])[None]).float().cuda() detection_boxes = detections[:, :self.cf.dim * 2 + 1] / scale with torch.no_grad(): detection_masks = self.mask(self.mrcnn_feature_maps, detection_boxes) return [rpn_pred_logits, rpn_pred_deltas, batch_proposal_boxes, detections, detection_masks] def loss_samples_forward(self, batch_gt_class_ids, batch_gt_boxes, batch_gt_masks): """ this is the second forward pass through the second stage (features from stage one are re-used). samples few rois in detection_target_layer and forwards only those for loss computation. :param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels. :param batch_gt_boxes: list over batch elements. Each element is a list over the corresponding roi target coordinates. :param batch_gt_masks: list over batch elements. Each element is binary mask of shape (n_gt_rois, y, x, (z), c) :return: sample_logits: (n_sampled_rois, n_classes) predicted class scores. :return: sample_boxes: (n_sampled_rois, n_classes, 2 * dim) predicted corrections to be applied to proposals for refinement. :return: sample_mask: (n_sampled_rois, n_classes, y, x, (z)) predicted masks per class and proposal. :return: sample_target_class_ids: (n_sampled_rois) target class labels of sampled proposals. :return: sample_target_deltas: (n_sampled_rois, 2 * dim) target deltas of sampled proposals for box refinement. :return: sample_target_masks: (n_sampled_rois, y, x, (z)) target masks of sampled proposals. :return: sample_proposals: (n_sampled_rois, 2 * dim) RPN output for sampled proposals. only for monitoring/plotting. """ # sample rois for loss and get corresponding targets for all Mask R-CNN head network losses. sample_ix, sample_target_class_ids, sample_target_deltas, sample_target_mask = \ detection_target_layer(self.rpn_rois_batch_info, self.batch_mrcnn_class_scores, batch_gt_class_ids, batch_gt_boxes, batch_gt_masks, self.cf) # re-use feature maps and RPN output from first forward pass. sample_proposals = self.rpn_rois_batch_info[sample_ix] if 0 not in sample_proposals.size(): sample_logits, sample_boxes = self.classifier(self.mrcnn_feature_maps, sample_proposals) sample_mask = self.mask(self.mrcnn_feature_maps, sample_proposals) else: sample_logits = torch.FloatTensor().cuda() sample_boxes = torch.FloatTensor().cuda() sample_mask = torch.FloatTensor().cuda() return [sample_logits, sample_boxes, sample_mask, sample_target_class_ids, sample_target_deltas, sample_target_mask, sample_proposals] \ No newline at end of file diff --git a/models/ufrcnn.py b/models/ufrcnn.py index 08c31bf..5f60e45 100644 --- a/models/ufrcnn.py +++ b/models/ufrcnn.py @@ -1,1273 +1,1273 @@ #!/usr/bin/env python # Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """ Parts are based on https://github.com/multimodallearning/pytorch-mask-rcnn published under MIT license. """ import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.utils sys.path.append("..") import utils.model_utils as mutils import utils.exp_utils as utils from custom_extensions.nms import nms from custom_extensions.roi_align import roi_align ############################################################ # Networks on top of backbone ############################################################ class RPN(nn.Module): """ Region Proposal Network. """ def __init__(self, cf, conv): super(RPN, self).__init__() self.dim = conv.dim self.conv_shared = conv(cf.end_filts, cf.n_rpn_features, ks=3, stride=cf.rpn_anchor_stride, pad=1, relu=cf.relu) self.conv_class = conv(cf.n_rpn_features, 2 * len(cf.rpn_anchor_ratios), ks=1, stride=1, relu=None) self.conv_bbox = conv(cf.n_rpn_features, 2 * self.dim * len(cf.rpn_anchor_ratios), ks=1, stride=1, relu=None) def forward(self, x): """ :param x: input feature maps (b, in_channels, y, x, (z)) :return: rpn_class_logits (b, 2, n_anchors) :return: rpn_probs_logits (b, 2, n_anchors) :return: rpn_bbox (b, 2 * dim, n_anchors) """ # Shared convolutional base of the RPN. x = self.conv_shared(x) # Anchor Score. (batch, anchors per location * 2, y, x, (z)). rpn_class_logits = self.conv_class(x) # Reshape to (batch, 2, anchors) axes = (0, 2, 3, 1) if self.dim == 2 else (0, 2, 3, 4, 1) rpn_class_logits = rpn_class_logits.permute(*axes) rpn_class_logits = rpn_class_logits.contiguous() rpn_class_logits = rpn_class_logits.view(x.size()[0], -1, 2) # Softmax on last dimension (fg vs. bg). rpn_probs = F.softmax(rpn_class_logits, dim=2) # Bounding box refinement. (batch, anchors_per_location * (y, x, (z), log(h), log(w), (log(d)), y, x, (z)) rpn_bbox = self.conv_bbox(x) # Reshape to (batch, 2*dim, anchors) rpn_bbox = rpn_bbox.permute(*axes) rpn_bbox = rpn_bbox.contiguous() rpn_bbox = rpn_bbox.view(x.size()[0], -1, self.dim * 2) return [rpn_class_logits, rpn_probs, rpn_bbox] class Classifier(nn.Module): """ Head network for classification and bounding box refinement. Performs RoiAlign, processes resulting features through a shared convolutional base and finally branches off the classifier- and regression head. """ def __init__(self, cf, conv): super(Classifier, self).__init__() self.dim = conv.dim self.in_channels = cf.end_filts self.pool_size = cf.pool_size self.pyramid_levels = cf.pyramid_levels # instance_norm does not work with spatial dims (1, 1, (1)) norm = cf.norm if cf.norm != 'instance_norm' else None self.conv1 = conv(cf.end_filts, cf.end_filts * 4, ks=self.pool_size, stride=1, norm=norm, relu=cf.relu) self.conv2 = conv(cf.end_filts * 4, cf.end_filts * 4, ks=1, stride=1, norm=norm, relu=cf.relu) self.linear_class = nn.Linear(cf.end_filts * 4, cf.head_classes) self.linear_bbox = nn.Linear(cf.end_filts * 4, cf.head_classes * 2 * self.dim) def forward(self, x, rois): """ :param x: input feature maps (b, in_channels, y, x, (z)) :param rois: normalized box coordinates as proposed by the RPN to be forwarded through the second stage (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix). Proposals of all batch elements have been merged to one vector, while the origin info has been stored for re-allocation. :return: mrcnn_class_logits (n_proposals, n_head_classes) :return: mrcnn_bbox (n_proposals, n_head_classes, 2 * dim) predicted corrections to be applied to proposals for refinement. """ x = pyramid_roi_align(x, rois, self.pool_size, self.pyramid_levels, self.dim) x = self.conv1(x) x = self.conv2(x) x = x.view(-1, self.in_channels * 4) mrcnn_class_logits = self.linear_class(x) mrcnn_bbox = self.linear_bbox(x) mrcnn_bbox = mrcnn_bbox.view(mrcnn_bbox.size()[0], -1, self.dim * 2) return [mrcnn_class_logits, mrcnn_bbox] class Mask(nn.Module): """ Head network for proposal-based mask segmentation. Performs RoiAlign, some convolutions and applies sigmoid on the output logits to allow for overlapping classes. """ def __init__(self, cf, conv): super(Mask, self).__init__() self.pool_size = cf.mask_pool_size self.pyramid_levels = cf.pyramid_levels self.dim = conv.dim self.conv1 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu) self.conv2 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu) self.conv3 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu) self.conv4 = conv(cf.end_filts, cf.end_filts, ks=3, stride=1, pad=1, norm=cf.norm, relu=cf.relu) if conv.dim == 2: self.deconv = nn.ConvTranspose2d(cf.end_filts, cf.end_filts, kernel_size=2, stride=2) else: self.deconv = nn.ConvTranspose3d(cf.end_filts, cf.end_filts, kernel_size=2, stride=2) self.relu = nn.ReLU(inplace=True) if cf.relu == 'relu' else nn.LeakyReLU(inplace=True) self.conv5 = conv(cf.end_filts, cf.head_classes, ks=1, stride=1, relu=None) self.sigmoid = nn.Sigmoid() def forward(self, x, rois): """ :param x: input feature maps (b, in_channels, y, x, (z)) :param rois: normalized box coordinates as proposed by the RPN to be forwarded through the second stage (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix). Proposals of all batch elements have been merged to one vector, while the origin info has been stored for re-allocation. :return: x: masks (n_sampled_proposals (n_detections in inference), n_classes, y, x, (z)) """ x = pyramid_roi_align(x, rois, self.pool_size, self.pyramid_levels, self.dim) x = self.conv1(x) x = self.conv2(x) x = self.conv3(x) x = self.conv4(x) x = self.relu(self.deconv(x)) x = self.conv5(x) x = self.sigmoid(x) return x ############################################################ # Loss Functions ############################################################ def compute_rpn_class_loss(rpn_match, rpn_class_logits, shem_poolsize): """ :param rpn_match: (n_anchors). [-1, 0, 1] for negative, neutral, and positive matched anchors. :param rpn_class_logits: (n_anchors, 2). logits from RPN classifier. :param shem_poolsize: int. factor of top-k candidates to draw from per negative sample (stochastic-hard-example-mining). :return: loss: torch tensor :return: np_neg_ix: 1D array containing indices of the neg_roi_logits, which have been sampled for training. """ # filter out neutral anchors. pos_indices = torch.nonzero(rpn_match == 1) neg_indices = torch.nonzero(rpn_match == -1) # loss for positive samples if 0 not in pos_indices.size(): pos_indices = pos_indices.squeeze(1) roi_logits_pos = rpn_class_logits[pos_indices] pos_loss = F.cross_entropy(roi_logits_pos, torch.LongTensor([1] * pos_indices.shape[0]).cuda()) else: pos_loss = torch.FloatTensor([0]).cuda() # loss for negative samples: draw hard negative examples (SHEM) # that match the number of positive samples, but at least 1. if 0 not in neg_indices.size(): neg_indices = neg_indices.squeeze(1) roi_logits_neg = rpn_class_logits[neg_indices] negative_count = np.max((1, pos_indices.cpu().data.numpy().size)) roi_probs_neg = F.softmax(roi_logits_neg, dim=1) neg_ix = mutils.shem(roi_probs_neg, negative_count, shem_poolsize) neg_loss = F.cross_entropy(roi_logits_neg[neg_ix], torch.LongTensor([0] * neg_ix.shape[0]).cuda()) np_neg_ix = neg_ix.cpu().data.numpy() else: neg_loss = torch.FloatTensor([0]).cuda() np_neg_ix = np.array([]).astype('int32') loss = (pos_loss + neg_loss) / 2 return loss, np_neg_ix def compute_rpn_bbox_loss(rpn_target_deltas, rpn_pred_deltas, rpn_match): """ :param rpn_target_deltas: (b, n_positive_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd)))). Uses 0 padding to fill in unsed bbox deltas. :param rpn_pred_deltas: predicted deltas from RPN. (b, n_anchors, (dy, dx, (dz), log(dh), log(dw), (log(dd)))) :param rpn_match: (n_anchors). [-1, 0, 1] for negative, neutral, and positive matched anchors. :return: loss: torch 1D tensor. """ if 0 not in torch.nonzero(rpn_match == 1).size(): indices = torch.nonzero(rpn_match == 1).squeeze(1) # Pick bbox deltas that contribute to the loss rpn_pred_deltas = rpn_pred_deltas[indices] # Trim target bounding box deltas to the same length as rpn_bbox. target_deltas = rpn_target_deltas[:rpn_pred_deltas.size()[0], :] # Smooth L1 loss loss = F.smooth_l1_loss(rpn_pred_deltas, target_deltas) else: loss = torch.FloatTensor([0]).cuda() return loss def compute_mrcnn_class_loss(target_class_ids, pred_class_logits): """ :param target_class_ids: (n_sampled_rois) batch dimension was merged into roi dimension. :param pred_class_logits: (n_sampled_rois, n_classes) :return: loss: torch 1D tensor. """ if 0 not in target_class_ids.size(): loss = F.cross_entropy(pred_class_logits, target_class_ids.long()) else: loss = torch.FloatTensor([0.]).cuda() return loss def compute_mrcnn_bbox_loss(mrcnn_target_deltas, mrcnn_pred_deltas, target_class_ids): """ :param mrcnn_target_deltas: (n_sampled_rois, (dy, dx, (dz), log(dh), log(dw), (log(dh))) :param mrcnn_pred_deltas: (n_sampled_rois, n_classes, (dy, dx, (dz), log(dh), log(dw), (log(dh))) :param target_class_ids: (n_sampled_rois) :return: loss: torch 1D tensor. """ if 0 not in torch.nonzero(target_class_ids > 0).size(): positive_roi_ix = torch.nonzero(target_class_ids > 0)[:, 0] positive_roi_class_ids = target_class_ids[positive_roi_ix].long() target_bbox = mrcnn_target_deltas[positive_roi_ix, :].detach() pred_bbox = mrcnn_pred_deltas[positive_roi_ix, positive_roi_class_ids, :] loss = F.smooth_l1_loss(pred_bbox, target_bbox) else: loss = torch.FloatTensor([0]).cuda() return loss def compute_mrcnn_mask_loss(target_masks, pred_masks, target_class_ids): """ :param target_masks: (n_sampled_rois, y, x, (z)) A float32 tensor of values 0 or 1. Uses zero padding to fill array. :param pred_masks: (n_sampled_rois, n_classes, y, x, (z)) float32 tensor with values between [0, 1]. :param target_class_ids: (n_sampled_rois) :return: loss: torch 1D tensor. """ if 0 not in torch.nonzero(target_class_ids > 0).size(): # Only positive ROIs contribute to the loss. And only # the class specific mask of each ROI. positive_ix = torch.nonzero(target_class_ids > 0)[:, 0] positive_class_ids = target_class_ids[positive_ix].long() y_true = target_masks[positive_ix, :, :].detach() y_pred = pred_masks[positive_ix, positive_class_ids, :, :] loss = F.binary_cross_entropy(y_pred, y_true) else: loss = torch.FloatTensor([0]).cuda() return loss ############################################################ # Helper Layers ############################################################ # def proposal_layer(rpn_pred_probs, rpn_pred_deltas, proposal_count, anchors, cf): # """ # Receives anchor scores and selects a subset to pass as proposals # to the second stage. Filtering is done based on anchor scores and # non-max suppression to remove overlaps. It also applies bounding # box refinment detals to anchors. # :param rpn_pred_probs: (b, n_anchors, 2) # :param rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d)))) # :return: batch_normalized_boxes: Proposals in normalized coordinates # (b, proposal_count, (y1, x1, y2, x2, (z1), (z2))) # :return: batch_out_proposals: Box coords + RPN foreground scores # for monitoring/plotting (b, proposal_count, (y1, x1, y2, x2, (z1), (z2), score)) # """ # batch_scores = rpn_pred_probs[:, :, 1] # batch_deltas = rpn_pred_deltas # batch_anchors = anchors # batch_normalized_boxes = [] # batch_out_proposals = [] # # # loop over batch dimension. # for ix in range(batch_scores.shape[0]): # # scores = batch_scores[ix] # deltas = batch_deltas[ix] # anchors = batch_anchors.clone() # # norm deltas # std_dev = torch.from_numpy(cf.rpn_bbox_std_dev[None]).float().cuda() # deltas = deltas * std_dev # # # improve performance by trimming to top anchors by score # # and doing the rest on the smaller subset. # pre_nms_limit = min(cf.pre_nms_limit, anchors.size()[0]) # scores, order = scores.sort(descending=True) # order = order[:pre_nms_limit] # scores = scores[:pre_nms_limit] # deltas = deltas[order, :] # anchors = anchors[order, :] # # # apply deltas to anchors to get refined anchors and filter with non-maximum surpression. # if batch_deltas.shape[-1] == 4: # boxes = mutils.apply_box_deltas_2D(anchors, deltas) # boxes = mutils.clip_boxes_2D(boxes, cf.window) # keep = nms_2D(torch.cat((boxes, scores.unsqueeze(1)), 1), cf.rpn_nms_threshold) # norm = torch.from_numpy(cf.scale).float().cuda() # # else: # boxes = mutils.apply_box_deltas_3D(anchors, deltas) # boxes = mutils.clip_boxes_3D(boxes, cf.window) # keep = nms_3D(torch.cat((boxes, scores.unsqueeze(1)), 1), cf.rpn_nms_threshold) # norm = torch.from_numpy(cf.scale).float().cuda() # # keep = keep[:proposal_count] # boxes = boxes[keep, :] # rpn_scores = scores[keep][:, None] # # # pad missing boxes with 0. # if boxes.shape[0] < proposal_count: # n_pad_boxes = proposal_count - boxes.shape[0] # zeros = torch.zeros([n_pad_boxes, boxes.shape[1]]).cuda() # boxes = torch.cat([boxes, zeros], dim=0) # zeros = torch.zeros([n_pad_boxes, rpn_scores.shape[1]]).cuda() # rpn_scores = torch.cat([rpn_scores, zeros], dim=0) # # # concat box and score info for monitoring/plotting. # batch_out_proposals.append(torch.cat((boxes, rpn_scores), 1).cpu().data.numpy()) # # normalize dimensions to range of 0 to 1. # normalized_boxes = boxes / norm # # add back batch dimension # batch_normalized_boxes.append(normalized_boxes.unsqueeze(0)) # # batch_normalized_boxes = torch.cat(batch_normalized_boxes) # batch_out_proposals = np.array(batch_out_proposals) # return batch_normalized_boxes, batch_out_proposals def refine_proposals(rpn_pred_probs, rpn_pred_deltas, proposal_count, batch_anchors, cf): """ Receives anchor scores and selects a subset to pass as proposals to the second stage. Filtering is done based on anchor scores and non-max suppression to remove overlaps. It also applies bounding box refinment details to anchors. :param rpn_pred_probs: (b, n_anchors, 2) :param rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d)))) :return: batch_normalized_props: Proposals in normalized coordinates (b, proposal_count, (y1, x1, y2, x2, (z1), (z2), score)) :return: batch_out_proposals: Box coords + RPN foreground scores for monitoring/plotting (b, proposal_count, (y1, x1, y2, x2, (z1), (z2), score)) """ std_dev = torch.from_numpy(cf.rpn_bbox_std_dev[None]).float().cuda() norm = torch.from_numpy(cf.scale).float().cuda() anchors = batch_anchors.clone() batch_scores = rpn_pred_probs[:, :, 1] # norm deltas batch_deltas = rpn_pred_deltas * std_dev batch_normalized_props = [] batch_out_proposals = [] # loop over batch dimension. for ix in range(batch_scores.shape[0]): scores = batch_scores[ix] deltas = batch_deltas[ix] # improve performance by trimming to top anchors by score # and doing the rest on the smaller subset. pre_nms_limit = min(cf.pre_nms_limit, anchors.size()[0]) scores, order = scores.sort(descending=True) order = order[:pre_nms_limit] scores = scores[:pre_nms_limit] deltas = deltas[order, :] # apply deltas to anchors to get refined anchors and filter with non-maximum suppression. if batch_deltas.shape[-1] == 4: boxes = mutils.apply_box_deltas_2D(anchors[order, :], deltas) boxes = mutils.clip_boxes_2D(boxes, cf.window) else: boxes = mutils.apply_box_deltas_3D(anchors[order, :], deltas) boxes = mutils.clip_boxes_3D(boxes, cf.window) # boxes are y1,x1,y2,x2, torchvision-nms requires x1,y1,x2,y2, but consistent swap x<->y is irrelevant. keep = nms.nms(boxes, scores, cf.rpn_nms_threshold) keep = keep[:proposal_count] boxes = boxes[keep, :] rpn_scores = scores[keep][:, None] # pad missing boxes with 0. if boxes.shape[0] < proposal_count: n_pad_boxes = proposal_count - boxes.shape[0] zeros = torch.zeros([n_pad_boxes, boxes.shape[1]]).cuda() boxes = torch.cat([boxes, zeros], dim=0) zeros = torch.zeros([n_pad_boxes, rpn_scores.shape[1]]).cuda() rpn_scores = torch.cat([rpn_scores, zeros], dim=0) # concat box and score info for monitoring/plotting. batch_out_proposals.append(torch.cat((boxes, rpn_scores), 1).cpu().data.numpy()) # normalize dimensions to range of 0 to 1. normalized_boxes = boxes / norm assert torch.all(normalized_boxes <= 1), "normalized box coords >1 found" # add again batch dimension batch_normalized_props.append(normalized_boxes.unsqueeze(0)) batch_normalized_props = torch.cat(batch_normalized_props) batch_out_proposals = np.array(batch_out_proposals) return batch_normalized_props, batch_out_proposals # def pyramid_roi_align(feature_maps, rois, pool_size, pyramid_levels, dim): # """ # Implements ROI Pooling on multiple levels of the feature pyramid. # :param feature_maps: list of feature maps, each of shape (b, c, y, x , (z)) # :param rois: proposals (normalized coords.) as returned by RPN. contain info about original batch element allocation. # (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ixs) # :param pool_size: list of poolsizes in dims: [x, y, (z)] # :param pyramid_levels: list. [0, 1, 2, ...] # :return: pooled: pooled feature map rois (n_proposals, c, poolsize_y, poolsize_x, (poolsize_z)) # # Output: # Pooled regions in the shape: [num_boxes, height, width, channels]. # The width and height are those specific in the pool_shape in the layer # constructor. # """ # boxes = rois[:, :dim*2] # batch_ixs = rois[:, dim*2] # # # Assign each ROI to a level in the pyramid based on the ROI area. # if dim == 2: # y1, x1, y2, x2 = boxes.chunk(4, dim=1) # else: # y1, x1, y2, x2, z1, z2 = boxes.chunk(6, dim=1) # # h = y2 - y1 # w = x2 - x1 # # # Equation 1 in https://arxiv.org/abs/1612.03144. Account for # # the fact that our coordinates are normalized here. # # divide sqrt(h*w) by 1 instead image_area. # roi_level = (4 + mutils.log2(torch.sqrt(h*w))).round().int().clamp(pyramid_levels[0], pyramid_levels[-1]) # # if Pyramid contains additional level P6, adapt the roi_level assignemnt accordingly. # if len(pyramid_levels) == 5: # roi_level[h*w > 0.65] = 5 # # # Loop through levels and apply ROI pooling to each. # pooled = [] # box_to_level = [] # for level_ix, level in enumerate(pyramid_levels): # ix = roi_level == level # if not ix.any(): # continue # ix = torch.nonzero(ix)[:, 0] # level_boxes = boxes[ix, :] # # re-assign rois to feature map of original batch element. # ind = batch_ixs[ix].int() # # # Keep track of which box is mapped to which level # box_to_level.append(ix) # # # Stop gradient propogation to ROI proposals # level_boxes = level_boxes.detach() # # # Crop and Resize # # From Mask R-CNN paper: "We sample four regular locations, so # # that we can evaluate either max or average pooling. In fact, # # interpolating only a single value at each bin center (without # # pooling) is nearly as effective." # # # # Here we use the simplified approach of a single value per bin, # # which is how is done in tf.crop_and_resize() # # # # Also fixed a bug from original implementation, reported in: # # https://hackernoon.com/how-tensorflows-tf-image-resize-stole-60-days-of-my-life-aba5eb093f35 # # if len(pool_size) == 2: # pooled_features = ra2D(pool_size[0], pool_size[1], 0)(feature_maps[level_ix], level_boxes, ind) # else: # pooled_features = ra3D(pool_size[0], pool_size[1], pool_size[2], 0)(feature_maps[level_ix], level_boxes, ind) # # pooled.append(pooled_features) # # # # Pack pooled features into one tensor # pooled = torch.cat(pooled, dim=0) # # # Pack box_to_level mapping into one array and add another # # column representing the order of pooled boxes # box_to_level = torch.cat(box_to_level, dim=0) # # # Rearrange pooled features to match the order of the original boxes # _, box_to_level = torch.sort(box_to_level) # pooled = pooled[box_to_level, :, :] # # return pooled def pyramid_roi_align(feature_maps, rois, pool_size, pyramid_levels, dim): """ Implements ROI Pooling on multiple levels of the feature pyramid. :param feature_maps: list of feature maps, each of shape (b, c, y, x , (z)) :param rois: proposals (normalized coords.) as returned by RPN. contain info about original batch element allocation. (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ixs) :param pool_size: list of poolsizes in dims: [x, y, (z)] :param pyramid_levels: list. [0, 1, 2, ...] :return: pooled: pooled feature map rois (n_proposals, c, poolsize_y, poolsize_x, (poolsize_z)) Output: Pooled regions in the shape: [num_boxes, height, width, channels]. The width and height are those specific in the pool_shape in the layer constructor. """ boxes = rois[:, :dim*2] batch_ixs = rois[:, dim*2] # Assign each ROI to a level in the pyramid based on the ROI area. if dim == 2: y1, x1, y2, x2 = boxes.chunk(4, dim=1) else: y1, x1, y2, x2, z1, z2 = boxes.chunk(6, dim=1) h = y2 - y1 w = x2 - x1 # Equation 1 in https://arxiv.org/abs/1612.03144. Account for # the fact that our coordinates are normalized here. # divide sqrt(h*w) by 1 instead image_area. roi_level = (4 + torch.log2(torch.sqrt(h*w))).round().int().clamp(pyramid_levels[0], pyramid_levels[-1]) # if Pyramid contains additional level P6, adapt the roi_level assignment accordingly. if len(pyramid_levels) == 5: roi_level[h*w > 0.65] = 5 # Loop through levels and apply ROI pooling to each. pooled = [] box_to_level = [] fmap_shapes = [f.shape for f in feature_maps] for level_ix, level in enumerate(pyramid_levels): ix = roi_level == level if not ix.any(): continue ix = torch.nonzero(ix)[:, 0] level_boxes = boxes[ix, :] # re-assign rois to feature map of original batch element. ind = batch_ixs[ix].int() # Keep track of which box is mapped to which level box_to_level.append(ix) # Stop gradient propogation to ROI proposals level_boxes = level_boxes.detach() if len(pool_size) == 2: # remap to feature map coordinate system y_exp, x_exp = fmap_shapes[level_ix][2:] # exp = expansion level_boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp], dtype=torch.float32).cuda()) pooled_features = roi_align.roi_align_2d(feature_maps[level_ix], torch.cat((ind.unsqueeze(1).float(), level_boxes), dim=1), pool_size) else: y_exp, x_exp, z_exp = fmap_shapes[level_ix][2:] level_boxes.mul_(torch.tensor([y_exp, x_exp, y_exp, x_exp, z_exp, z_exp], dtype=torch.float32).cuda()) pooled_features = roi_align.roi_align_3d(feature_maps[level_ix], torch.cat((ind.unsqueeze(1).float(), level_boxes), dim=1), pool_size) pooled.append(pooled_features) # Pack pooled features into one tensor pooled = torch.cat(pooled, dim=0) # Pack box_to_level mapping into one array and add another # column representing the order of pooled boxes box_to_level = torch.cat(box_to_level, dim=0) # Rearrange pooled features to match the order of the original boxes _, box_to_level = torch.sort(box_to_level) pooled = pooled[box_to_level, :, :] return pooled def detection_target_layer(batch_proposals, batch_mrcnn_class_scores, batch_gt_class_ids, batch_gt_boxes, cf): """ Subsamples proposals for mrcnn losses and generates targets. Sampling is done per batch element, seems to have positive effects on training, as opposed to sampling over entire batch. Negatives are sampled via stochastic-hard-example-mining (SHEM), where a number of negative proposals are drawn from larger pool of highest scoring proposals for stochasticity. Scoring is obtained here as the max over all foreground probabilities as returned by mrcnn_classifier (worked better than loss-based class balancing methods like "online-hard-example-mining" or "focal loss".) :param batch_proposals: (n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ixs). boxes as proposed by RPN. n_proposals here is determined by batch_size * POST_NMS_ROIS. :param batch_mrcnn_class_scores: (n_proposals, n_classes) :param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels. :param batch_gt_boxes: list over batch elements. Each element is a list over the corresponding roi target coordinates. :param batch_gt_masks: list over batch elements. Each element is binary mask of shape (n_gt_rois, y, x, (z), c) :return: sample_indices: (n_sampled_rois) indices of sampled proposals to be used for loss functions. :return: target_class_ids: (n_sampled_rois)containing target class labels of sampled proposals. :return: target_deltas: (n_sampled_rois, 2 * dim) containing target deltas of sampled proposals for box refinement. :return: target_masks: (n_sampled_rois, y, x, (z)) containing target masks of sampled proposals. """ # normalization of target coordinates if cf.dim == 2: h, w = cf.patch_size scale = torch.from_numpy(np.array([h, w, h, w])).float().cuda() else: h, w, z = cf.patch_size scale = torch.from_numpy(np.array([h, w, h, w, z, z])).float().cuda() positive_count = 0 negative_count = 0 sample_positive_indices = [] sample_negative_indices = [] sample_deltas = [] sample_class_ids = [] std_dev = torch.from_numpy(cf.bbox_std_dev).float().cuda() # loop over batch and get positive and negative sample rois. for b in range(len(batch_gt_class_ids)): gt_class_ids = torch.from_numpy(batch_gt_class_ids[b]).int().cuda() if np.any(batch_gt_class_ids[b] > 0): # skip roi selection for no gt images. gt_boxes = torch.from_numpy(batch_gt_boxes[b]).float().cuda() / scale else: gt_boxes = torch.FloatTensor().cuda() # get proposals and indices of current batch element. proposals = batch_proposals[batch_proposals[:, -1] == b][:, :-1] batch_element_indices = torch.nonzero(batch_proposals[:, -1] == b).squeeze(1) # Compute overlaps matrix [proposals, gt_boxes] if 0 not in gt_boxes.size(): if gt_boxes.shape[1] == 4: overlaps = mutils.bbox_overlaps_2D(proposals, gt_boxes) else: overlaps = mutils.bbox_overlaps_3D(proposals, gt_boxes) # Determine postive and negative ROIs roi_iou_max = torch.max(overlaps, dim=1)[0] # 1. Positive ROIs are those with >= 0.5 IoU with a GT box positive_roi_bool = roi_iou_max >= (0.5 if cf.dim == 2 else 0.3) # 2. Negative ROIs are those with < 0.1 with every GT box. negative_roi_bool = roi_iou_max < (0.1 if cf.dim == 2 else 0.01) else: positive_roi_bool = torch.FloatTensor().cuda() negative_roi_bool = torch.from_numpy(np.array([1]*proposals.shape[0])).cuda() # Sample Positive ROIs if 0 not in torch.nonzero(positive_roi_bool).size(): positive_indices = torch.nonzero(positive_roi_bool).squeeze(1) positive_samples = int(cf.train_rois_per_image * cf.roi_positive_ratio) rand_idx = torch.randperm(positive_indices.size()[0]) rand_idx = rand_idx[:positive_samples].cuda() positive_indices = positive_indices[rand_idx] positive_samples = positive_indices.size()[0] positive_rois = proposals[positive_indices, :] # Assign positive ROIs to GT boxes. positive_overlaps = overlaps[positive_indices, :] roi_gt_box_assignment = torch.max(positive_overlaps, dim=1)[1] roi_gt_boxes = gt_boxes[roi_gt_box_assignment, :] roi_gt_class_ids = gt_class_ids[roi_gt_box_assignment] # Compute bbox refinement targets for positive ROIs deltas = mutils.box_refinement(positive_rois, roi_gt_boxes) deltas /= std_dev sample_positive_indices.append(batch_element_indices[positive_indices]) sample_deltas.append(deltas) sample_class_ids.append(roi_gt_class_ids) positive_count += positive_samples else: positive_samples = 0 # Negative ROIs. Add enough to maintain positive:negative ratio, but at least 1. Sample via SHEM. if 0 not in torch.nonzero(negative_roi_bool).size(): negative_indices = torch.nonzero(negative_roi_bool).squeeze(1) r = 1.0 / cf.roi_positive_ratio b_neg_count = np.max((int(r * positive_samples - positive_samples), 1)) roi_probs_neg = batch_mrcnn_class_scores[batch_element_indices[negative_indices]] raw_sampled_indices = mutils.shem(roi_probs_neg, b_neg_count, cf.shem_poolsize) sample_negative_indices.append(batch_element_indices[negative_indices[raw_sampled_indices]]) negative_count += raw_sampled_indices.size()[0] if len(sample_positive_indices) > 0: target_deltas = torch.cat(sample_deltas) target_class_ids = torch.cat(sample_class_ids) # Pad target information with zeros for negative ROIs. if positive_count > 0 and negative_count > 0: sample_indices = torch.cat((torch.cat(sample_positive_indices), torch.cat(sample_negative_indices)), dim=0) zeros = torch.zeros(negative_count).int().cuda() target_class_ids = torch.cat([target_class_ids, zeros], dim=0) zeros = torch.zeros(negative_count, cf.dim * 2).cuda() target_deltas = torch.cat([target_deltas, zeros], dim=0) elif positive_count > 0: sample_indices = torch.cat(sample_positive_indices) elif negative_count > 0: sample_indices = torch.cat(sample_negative_indices) zeros = torch.zeros(negative_count).int().cuda() target_class_ids = zeros zeros = torch.zeros(negative_count, cf.dim * 2).cuda() target_deltas = zeros else: sample_indices = torch.LongTensor().cuda() target_class_ids = torch.IntTensor().cuda() target_deltas = torch.FloatTensor().cuda() return sample_indices, target_class_ids, target_deltas ############################################################ # Output Handler ############################################################ # def refine_detections(rois, probs, deltas, batch_ixs, cf): # """ # Refine classified proposals, filter overlaps and return final detections. # # :param rois: (n_proposals, 2 * dim) normalized boxes as proposed by RPN. n_proposals = batch_size * POST_NMS_ROIS # :param probs: (n_proposals, n_classes) softmax probabilities for all rois as predicted by mrcnn classifier. # :param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by mrcnn bbox regressor. # :param batch_ixs: (n_proposals) batch element assignemnt info for re-allocation. # :return: result: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score)) # """ # # class IDs per ROI. Since scores of all classes are of interest (not just max class), all are kept at this point. # class_ids = [] # fg_classes = cf.head_classes - 1 # # repeat vectors to fill in predictions for all foreground classes. # for ii in range(1, fg_classes + 1): # class_ids += [ii] * rois.shape[0] # class_ids = torch.from_numpy(np.array(class_ids)).cuda() # # rois = rois.repeat(fg_classes, 1) # probs = probs.repeat(fg_classes, 1) # deltas = deltas.repeat(fg_classes, 1, 1) # batch_ixs = batch_ixs.repeat(fg_classes) # # # get class-specific scores and bounding box deltas # idx = torch.arange(class_ids.size()[0]).long().cuda() # class_scores = probs[idx, class_ids] # deltas_specific = deltas[idx, class_ids] # batch_ixs = batch_ixs[idx] # # # apply bounding box deltas. re-scale to image coordinates. # std_dev = torch.from_numpy(np.reshape(cf.rpn_bbox_std_dev, [1, cf.dim * 2])).float().cuda() # scale = torch.from_numpy(cf.scale).float().cuda() # refined_rois = mutils.apply_box_deltas_2D(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \ # mutils.apply_box_deltas_3D(rois, deltas_specific * std_dev) * scale # # # round and cast to int since we're deadling with pixels now # refined_rois = mutils.clip_to_window(cf.window, refined_rois) # refined_rois = torch.round(refined_rois) # # # filter out low confidence boxes # keep = idx # keep_bool = (class_scores >= cf.model_min_confidence) # if 0 not in torch.nonzero(keep_bool).size(): # # score_keep = torch.nonzero(keep_bool)[:, 0] # pre_nms_class_ids = class_ids[score_keep] # pre_nms_rois = refined_rois[score_keep] # pre_nms_scores = class_scores[score_keep] # pre_nms_batch_ixs = batch_ixs[score_keep] # # for j, b in enumerate(mutils.unique1d(pre_nms_batch_ixs)): # # bixs = torch.nonzero(pre_nms_batch_ixs == b)[:, 0] # bix_class_ids = pre_nms_class_ids[bixs] # bix_rois = pre_nms_rois[bixs] # bix_scores = pre_nms_scores[bixs] # # for i, class_id in enumerate(mutils.unique1d(bix_class_ids)): # # ixs = torch.nonzero(bix_class_ids == class_id)[:, 0] # # nms expects boxes sorted by score. # ix_rois = bix_rois[ixs] # ix_scores = bix_scores[ixs] # ix_scores, order = ix_scores.sort(descending=True) # ix_rois = ix_rois[order, :] # # if cf.dim == 2: # class_keep = nms_2D(torch.cat((ix_rois, ix_scores.unsqueeze(1)), dim=1), cf.detection_nms_threshold) # else: # class_keep = nms_3D(torch.cat((ix_rois, ix_scores.unsqueeze(1)), dim=1), cf.detection_nms_threshold) # # # map indices back. # class_keep = keep[score_keep[bixs[ixs[order[class_keep]]]]] # # merge indices over classes for current batch element # b_keep = class_keep if i == 0 else mutils.unique1d(torch.cat((b_keep, class_keep))) # # # only keep top-k boxes of current batch-element # top_ids = class_scores[b_keep].sort(descending=True)[1][:cf.model_max_instances_per_batch_element] # b_keep = b_keep[top_ids] # # # merge indices over batch elements. # batch_keep = b_keep if j == 0 else mutils.unique1d(torch.cat((batch_keep, b_keep))) # # keep = batch_keep # # else: # keep = torch.tensor([0]).long().cuda() # # # arrange output # result = torch.cat((refined_rois[keep], # batch_ixs[keep].unsqueeze(1), # class_ids[keep].unsqueeze(1).float(), # class_scores[keep].unsqueeze(1)), dim=1) # # return result def refine_detections(cf, batch_ixs, rois, deltas, scores): """ Refine classified proposals (apply deltas to rpn rois), filter overlaps (nms) and return final detections. :param rois: (n_proposals, 2 * dim) normalized boxes as proposed by RPN. n_proposals = batch_size * POST_NMS_ROIS :param deltas: (n_proposals, n_classes, 2 * dim) box refinement deltas as predicted by mrcnn bbox regressor. :param batch_ixs: (n_proposals) batch element assignment info for re-allocation. :param scores: (n_proposals, n_classes) probabilities for all classes per roi as predicted by mrcnn classifier. :return: result: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score, *regression vector features)) """ # class IDs per ROI. Since scores of all classes are of interest (not just max class), all are kept at this point. class_ids = [] fg_classes = cf.head_classes - 1 # repeat vectors to fill in predictions for all foreground classes. for ii in range(1, fg_classes + 1): class_ids += [ii] * rois.shape[0] class_ids = torch.from_numpy(np.array(class_ids)).cuda() batch_ixs = batch_ixs.repeat(fg_classes) rois = rois.repeat(fg_classes, 1) deltas = deltas.repeat(fg_classes, 1, 1) scores = scores.repeat(fg_classes, 1) # get class-specific scores and bounding box deltas idx = torch.arange(class_ids.size()[0]).long().cuda() # using idx instead of slice [:,] squashes first dimension. #len(class_ids)>scores.shape[1] --> probs is broadcasted by expansion from fg_classes-->len(class_ids) batch_ixs = batch_ixs[idx] deltas_specific = deltas[idx, class_ids] class_scores = scores[idx, class_ids] # apply bounding box deltas. re-scale to image coordinates. std_dev = torch.from_numpy(np.reshape(cf.rpn_bbox_std_dev, [1, cf.dim * 2])).float().cuda() scale = torch.from_numpy(cf.scale).float().cuda() refined_rois = mutils.apply_box_deltas_2D(rois, deltas_specific * std_dev) * scale if cf.dim == 2 else \ mutils.apply_box_deltas_3D(rois, deltas_specific * std_dev) * scale # round and cast to int since we're dealing with pixels now refined_rois = mutils.clip_to_window(cf.window, refined_rois) refined_rois = torch.round(refined_rois) # filter out low confidence boxes keep = idx keep_bool = (class_scores >= cf.model_min_confidence) if not 0 in torch.nonzero(keep_bool).size(): score_keep = torch.nonzero(keep_bool)[:, 0] pre_nms_class_ids = class_ids[score_keep] pre_nms_rois = refined_rois[score_keep] pre_nms_scores = class_scores[score_keep] pre_nms_batch_ixs = batch_ixs[score_keep] for j, b in enumerate(mutils.unique1d(pre_nms_batch_ixs)): bixs = torch.nonzero(pre_nms_batch_ixs == b)[:, 0] bix_class_ids = pre_nms_class_ids[bixs] bix_rois = pre_nms_rois[bixs] bix_scores = pre_nms_scores[bixs] for i, class_id in enumerate(mutils.unique1d(bix_class_ids)): ixs = torch.nonzero(bix_class_ids == class_id)[:, 0] # nms expects boxes sorted by score. ix_rois = bix_rois[ixs] ix_scores = bix_scores[ixs] ix_scores, order = ix_scores.sort(descending=True) ix_rois = ix_rois[order, :] class_keep = nms.nms(ix_rois, ix_scores, cf.detection_nms_threshold) # map indices back. class_keep = keep[score_keep[bixs[ixs[order[class_keep]]]]] # merge indices over classes for current batch element b_keep = class_keep if i == 0 else mutils.unique1d(torch.cat((b_keep, class_keep))) # only keep top-k boxes of current batch-element top_ids = class_scores[b_keep].sort(descending=True)[1][:cf.model_max_instances_per_batch_element] b_keep = b_keep[top_ids] # merge indices over batch elements. batch_keep = b_keep if j == 0 else mutils.unique1d(torch.cat((batch_keep, b_keep))) keep = batch_keep else: keep = torch.tensor([0]).long().cuda() # arrange output output = [refined_rois[keep], batch_ixs[keep].unsqueeze(1)] output += [class_ids[keep].unsqueeze(1).float(), class_scores[keep].unsqueeze(1)] result = torch.cat(output, dim=1) # shape: (n_keeps, catted feats), catted feats: [0:dim*2] are box_coords, [dim*2] are batch_ics, # [dim*2+1] are class_ids, [dim*2+2] are scores, [dim*2+3:] are regression vector features (incl uncertainty) return result def get_results(cf, img_shape, detections, seg_logits, box_results_list=None): """ Restores batch dimension of merged detections, unmolds detections, creates and fills results dict. :param img_shape: :param detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score) :param detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head. :param box_results_list: None or list of output boxes for monitoring/plotting. each element is a list of boxes per batch element. :param return_masks: boolean. If True, full resolution masks are returned for all proposals (speed trade-off). :return: results_dict: dictionary with keys: 'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary: [[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...] 'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, 1] only fg. vs. bg for now. class-specific return of masks will come with implementation of instance segmentation evaluation. """ detections = detections.cpu().data.numpy() # restore batch dimension of merged detections using the batch_ix info. batch_ixs = detections[:, cf.dim*2] detections = [detections[batch_ixs == ix] for ix in range(img_shape[0])] # for test_forward, where no previous list exists. if box_results_list is None: box_results_list = [[] for _ in range(img_shape[0])] seg_preds = [] # loop over batch and unmold detections. for ix in range(img_shape[0]): if 0 not in detections[ix].shape: boxes = detections[ix][:, :2 * cf.dim].astype(np.int32) class_ids = detections[ix][:, 2 * cf.dim + 1].astype(np.int32) scores = detections[ix][:, 2 * cf.dim + 2] # Filter out detections with zero area. Often only happens in early # stages of training when the network weights are still a bit random. if cf.dim == 2: exclude_ix = np.where((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) <= 0)[0] else: exclude_ix = np.where( (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 5] - boxes[:, 4]) <= 0)[0] if exclude_ix.shape[0] > 0: boxes = np.delete(boxes, exclude_ix, axis=0) class_ids = np.delete(class_ids, exclude_ix, axis=0) scores = np.delete(scores, exclude_ix, axis=0) # add final perdictions to results. if 0 not in boxes.shape: for ix2, score in enumerate(scores): if score >= cf.model_min_confidence: box_results_list[ix].append({'box_coords': boxes[ix2], 'box_score': score, 'box_type': 'det', 'box_pred_class_id': class_ids[ix2]}) # create and fill results dictionary. results_dict = {'boxes': box_results_list} if seg_logits is None: # output dummy segmentation for retina_net. results_dict['seg_preds'] = np.zeros(img_shape)[:, 0][:, np.newaxis] else: # output label maps for retina_unet. results_dict['seg_preds'] = F.softmax(seg_logits, 1).argmax(1).cpu().data.numpy()[:, np.newaxis].astype('uint8') return results_dict ############################################################ # Mask R-CNN Class ############################################################ class net(nn.Module): def __init__(self, cf, logger): super(net, self).__init__() self.cf = cf self.logger = logger self.build() if self.cf.weight_init is not None: logger.info("using pytorch weight init of type {}".format(self.cf.weight_init)) mutils.initialize_weights(self) else: logger.info("using default pytorch weight init") def build(self): """Build Mask R-CNN architecture.""" # Image size must be dividable by 2 multiple times. h, w = self.cf.patch_size[:2] if h / 2**5 != int(h / 2**5) or w / 2**5 != int(w / 2**5): raise Exception("Image size must be dividable by 2 at least 5 times " "to avoid fractions when downscaling and upscaling." "For example, use 256, 320, 384, 448, 512, ... etc. ") # instanciate abstract multi dimensional conv class and backbone class. conv = mutils.NDConvGenerator(self.cf.dim) backbone = utils.import_module('bbone', self.cf.backbone_path) # build Anchors, FPN, RPN, Classifier / Bbox-Regressor -head, Mask-head self.np_anchors = mutils.generate_pyramid_anchors(self.logger, self.cf) self.anchors = torch.from_numpy(self.np_anchors).float().cuda() self.fpn = backbone.FPN(self.cf, conv, operate_stride1=True) self.rpn = RPN(self.cf, conv) self.classifier = Classifier(self.cf, conv) self.mask = Mask(self.cf, conv) self.final_conv = conv(self.cf.end_filts, self.cf.num_seg_classes, ks=1, pad=0, norm=self.cf.norm, relu=None) def train_forward(self, batch, is_validation=False): """ train method (also used for validation monitoring). wrapper around forward pass of network. prepares input data for processing, computes losses, and stores outputs in a dictionary. :param batch: dictionary containing 'data', 'seg', etc. :return: results_dict: dictionary with keys: 'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary: [[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...] 'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, n_classes]. 'torch_loss': 1D torch tensor for backprop. 'class_loss': classification loss for monitoring. """ img = batch['data'] gt_class_ids = batch['roi_labels'] gt_boxes = batch['bb_target'] axes = (0, 2, 3, 1) if self.cf.dim == 2 else (0, 2, 3, 4, 1) var_seg_ohe = torch.FloatTensor(mutils.get_one_hot_encoding(batch['seg'], self.cf.num_seg_classes)).cuda() var_seg = torch.LongTensor(batch['seg']).cuda() img = torch.from_numpy(img).float().cuda() batch_rpn_class_loss = torch.FloatTensor([0]).cuda() batch_rpn_bbox_loss = torch.FloatTensor([0]).cuda() # list of output boxes for monitoring/plotting. each element is a list of boxes per batch element. box_results_list = [[] for _ in range(img.shape[0])] #forward passes. 1. general forward pass, where no activations are saved in second stage (for performance # monitoring and loss sampling). 2. second stage forward pass of sampled rois with stored activations for backprop. rpn_class_logits, rpn_pred_deltas, proposal_boxes, detections, seg_logits = self.forward(img) mrcnn_class_logits, mrcnn_pred_deltas, target_class_ids, mrcnn_target_deltas, \ sample_proposals = self.loss_samples_forward(gt_class_ids, gt_boxes) # loop over batch for b in range(img.shape[0]): if len(gt_boxes[b]) > 0: # add gt boxes to output list for monitoring. for ix in range(len(gt_boxes[b])): box_results_list[b].append({'box_coords': batch['bb_target'][b][ix], 'box_label': batch['roi_labels'][b][ix], 'box_type': 'gt'}) # match gt boxes with anchors to generate targets for RPN losses. rpn_match, rpn_target_deltas = mutils.gt_anchor_matching(self.cf, self.np_anchors, gt_boxes[b]) # add positive anchors used for loss to output list for monitoring. pos_anchors = mutils.clip_boxes_numpy(self.np_anchors[np.argwhere(rpn_match == 1)][:, 0], img.shape[2:]) for p in pos_anchors: box_results_list[b].append({'box_coords': p, 'box_type': 'pos_anchor'}) else: rpn_match = np.array([-1]*self.np_anchors.shape[0]) rpn_target_deltas = np.array([0]) rpn_match_gpu = torch.from_numpy(rpn_match).cuda() rpn_target_deltas = torch.from_numpy(rpn_target_deltas).float().cuda() # compute RPN losses. rpn_class_loss, neg_anchor_ix = compute_rpn_class_loss(rpn_match_gpu, rpn_class_logits[b], self.cf.shem_poolsize) rpn_bbox_loss = compute_rpn_bbox_loss(rpn_target_deltas, rpn_pred_deltas[b], rpn_match_gpu) batch_rpn_class_loss += rpn_class_loss / img.shape[0] batch_rpn_bbox_loss += rpn_bbox_loss / img.shape[0] # add negative anchors used for loss to output list for monitoring. neg_anchors = mutils.clip_boxes_numpy(self.np_anchors[rpn_match == -1][neg_anchor_ix], img.shape[2:]) for n in neg_anchors: box_results_list[b].append({'box_coords': n, 'box_type': 'neg_anchor'}) # add highest scoring proposals to output list for monitoring. rpn_proposals = proposal_boxes[b][proposal_boxes[b, :, -1].argsort()][::-1] for r in rpn_proposals[:self.cf.n_plot_rpn_props, :-1]: box_results_list[b].append({'box_coords': r, 'box_type': 'prop'}) # add positive and negative roi samples used for mrcnn losses to output list for monitoring. if 0 not in sample_proposals.shape: rois = mutils.clip_to_window(self.cf.window, sample_proposals).cpu().data.numpy() for ix, r in enumerate(rois): box_results_list[int(r[-1])].append({'box_coords': r[:-1] * self.cf.scale, 'box_type': 'pos_class' if target_class_ids[ix] > 0 else 'neg_class'}) batch_rpn_class_loss = batch_rpn_class_loss batch_rpn_bbox_loss = batch_rpn_bbox_loss # compute mrcnn losses. mrcnn_class_loss = compute_mrcnn_class_loss(target_class_ids, mrcnn_class_logits) mrcnn_bbox_loss = compute_mrcnn_bbox_loss(mrcnn_target_deltas, mrcnn_pred_deltas, target_class_ids) # mrcnn can be run without pixelwise annotations available (Faster R-CNN mode). # In this case, the mask_loss is taken out of training. # if not self.cf.frcnn_mode: # mrcnn_mask_loss = compute_mrcnn_mask_loss(target_mask, mrcnn_pred_mask, target_class_ids) # else: # mrcnn_mask_loss = torch.FloatTensor([0]).cuda() seg_loss_dice = 1 - mutils.batch_dice(F.softmax(seg_logits, dim=1), var_seg_ohe) seg_loss_ce = F.cross_entropy(seg_logits, var_seg[:, 0]) loss = batch_rpn_class_loss + batch_rpn_bbox_loss + mrcnn_class_loss + mrcnn_bbox_loss + (seg_loss_dice + seg_loss_ce) / 2 # monitor RPN performance: detection count = the number of correctly matched proposals per fg-class. dcount = [list(target_class_ids.cpu().data.numpy()).count(c) for c in np.arange(self.cf.head_classes)[1:]] # run unmolding of predictions for monitoring and merge all results to one dictionary. results_dict = get_results(self.cf, img.shape, detections, seg_logits, box_results_list) results_dict['torch_loss'] = loss results_dict['monitor_values'] = {'loss': loss.item(), 'class_loss': mrcnn_class_loss.item()} results_dict['logger_string'] = "loss: {0:.2f}, rpn_class: {1:.2f}, rpn_bbox: {2:.2f}, mrcnn_class: {3:.2f}, " \ "mrcnn_bbox: {4:.2f}, dice_loss: {5:.2f}, dcount {6}"\ .format(loss.item(), batch_rpn_class_loss.item(), batch_rpn_bbox_loss.item(), mrcnn_class_loss.item(), mrcnn_bbox_loss.item(), seg_loss_dice.item(), dcount) return results_dict def test_forward(self, batch, return_masks=True): """ test method. wrapper around forward pass of network without usage of any ground truth information. prepares input data for processing and stores outputs in a dictionary. :param batch: dictionary containing 'data' :param return_masks: boolean. If True, full resolution masks are returned for all proposals (speed trade-off). :return: results_dict: dictionary with keys: 'boxes': list over batch elements. each batch element is a list of boxes. each box is a dictionary: [[{box_0}, ... {box_n}], [{box_0}, ... {box_n}], ...] 'seg_preds': pixel-wise class predictions (b, 1, y, x, (z)) with values [0, n_classes] """ img = batch['data'] img = torch.from_numpy(img).float().cuda() _, _, _, detections, seg_logits = self.forward(img) results_dict = get_results(self.cf, img.shape, detections, seg_logits) return results_dict def forward(self, img, is_training=True): """ :param img: input images (b, c, y, x, (z)). :return: rpn_pred_logits: (b, n_anchors, 2) :return: rpn_pred_deltas: (b, n_anchors, (y, x, (z), log(h), log(w), (log(d)))) :return: batch_proposal_boxes: (b, n_proposals, (y1, x1, y2, x2, (z1), (z2), batch_ix)) only for monitoring/plotting. :return: detections: (n_final_detections, (y1, x1, y2, x2, (z1), (z2), batch_ix, pred_class_id, pred_score) :return: detection_masks: (n_final_detections, n_classes, y, x, (z)) raw molded masks as returned by mask-head. """ # extract features. fpn_outs = self.fpn(img) seg_logits = self.final_conv(fpn_outs[0]) rpn_feature_maps = [fpn_outs[i + 1] for i in self.cf.pyramid_levels] self.mrcnn_feature_maps = rpn_feature_maps # loop through pyramid layers and apply RPN. layer_outputs = [] # list of lists for p in rpn_feature_maps: layer_outputs.append(self.rpn(p)) # concatenate layer outputs. # convert from list of lists of level outputs to list of lists of outputs across levels. # e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]] outputs = list(zip(*layer_outputs)) outputs = [torch.cat(list(o), dim=1) for o in outputs] rpn_pred_logits, rpn_pred_probs, rpn_pred_deltas = outputs # generate proposals: apply predicted deltas to anchors and filter by foreground scores from RPN classifier. proposal_count = self.cf.post_nms_rois_training if is_training else self.cf.post_nms_rois_inference batch_rpn_rois, batch_proposal_boxes = refine_proposals(rpn_pred_probs, rpn_pred_deltas, proposal_count, self.anchors, self.cf) # merge batch dimension of proposals while storing allocation info in coordinate dimension. batch_ixs = torch.from_numpy(np.repeat(np.arange(batch_rpn_rois.shape[0]), batch_rpn_rois.shape[1])).float().cuda() rpn_rois = batch_rpn_rois.view(-1, batch_rpn_rois.shape[2]) self.rpn_rois_batch_info = torch.cat((rpn_rois, batch_ixs.unsqueeze(1)), dim=1) # this is the first of two forward passes in the second stage, where no activations are stored for backprop. # here, all proposals are forwarded (with virtual_batch_size = batch_size * post_nms_rois.) # for inference/monitoring as well as sampling of rois for the loss functions. # processed in chunks of roi_chunk_size to re-adjust to gpu-memory. chunked_rpn_rois = self.rpn_rois_batch_info.split(self.cf.roi_chunk_size) class_logits_list, bboxes_list = [], [] with torch.no_grad(): for chunk in chunked_rpn_rois: chunk_class_logits, chunk_bboxes = self.classifier(self.mrcnn_feature_maps, chunk) class_logits_list.append(chunk_class_logits) bboxes_list.append(chunk_bboxes) batch_mrcnn_class_logits = torch.cat(class_logits_list, 0) batch_mrcnn_bbox = torch.cat(bboxes_list, 0) self.batch_mrcnn_class_scores = F.softmax(batch_mrcnn_class_logits, dim=1) # refine classified proposals, filter and return final detections. - detections = refine_detections(self.cf, batch_ixs, rpn_rois, batch_mrcnn_bbox, batch_ixs) + detections = refine_detections(self.cf, batch_ixs, rpn_rois, batch_mrcnn_bbox, self.batch_mrcnn_class_scores) return [rpn_pred_logits, rpn_pred_deltas, batch_proposal_boxes, detections, seg_logits] def loss_samples_forward(self, batch_gt_class_ids, batch_gt_boxes): """ this is the second forward pass through the second stage (features from stage one are re-used). samples few rois in detection_target_layer and forwards only those for loss computation. :param batch_gt_class_ids: list over batch elements. Each element is a list over the corresponding roi target labels. :param batch_gt_boxes: list over batch elements. Each element is a list over the corresponding roi target coordinates. :param batch_gt_masks: list over batch elements. Each element is binary mask of shape (n_gt_rois, y, x, (z), c) :return: sample_logits: (n_sampled_rois, n_classes) predicted class scores. :return: sample_boxes: (n_sampled_rois, n_classes, 2 * dim) predicted corrections to be applied to proposals for refinement. :return: sample_mask: (n_sampled_rois, n_classes, y, x, (z)) predicted masks per class and proposal. :return: sample_target_class_ids: (n_sampled_rois) target class labels of sampled proposals. :return: sample_target_deltas: (n_sampled_rois, 2 * dim) target deltas of sampled proposals for box refinement. :return: sample_target_masks: (n_sampled_rois, y, x, (z)) target masks of sampled proposals. :return: sample_proposals: (n_sampled_rois, 2 * dim) RPN output for sampled proposals. only for monitoring/plotting. """ # sample rois for loss and get corresponding targets for all Mask R-CNN head network losses. sample_ix, sample_target_class_ids, sample_target_deltas = \ detection_target_layer(self.rpn_rois_batch_info, self.batch_mrcnn_class_scores, batch_gt_class_ids, batch_gt_boxes, self.cf) # re-use feature maps and RPN output from first forward pass. sample_proposals = self.rpn_rois_batch_info[sample_ix] if 0 not in sample_proposals.size(): sample_logits, sample_boxes = self.classifier(self.mrcnn_feature_maps, sample_proposals) else: sample_logits = torch.FloatTensor().cuda() sample_boxes = torch.FloatTensor().cuda() return [sample_logits, sample_boxes, sample_target_class_ids, sample_target_deltas, sample_proposals] diff --git a/requirements.txt b/requirements.txt index eded0cd..8929526 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,27 +1,42 @@ +backcall==0.1.0 batchgenerators==0.19.3 cffi==1.11.5 cycler==0.10.0 +Cython==0.29.14 decorator==4.4.1 future==0.18.2 imageio==2.6.1 +ipython==7.10.2 +ipython-genutils==0.2.0 +jedi==0.15.1 kiwisolver==1.1.0 linecache2==1.0.0 matplotlib==3.1.2 networkx==2.4 -numpy==1.15.3 -pandas==0.23.4 +numpy==1.17.4 +pandas==0.25.3 +parso==0.5.2 +pexpect==4.7.0 +pickleshare==0.7.5 Pillow==6.2.1 +prompt-toolkit==3.0.2 +ptyprocess==0.6.0 pycparser==2.19 +Pygments==2.5.2 pyparsing==2.4.5 python-dateutil==2.8.1 pytz==2019.3 PyWavelets==1.1.1 scikit-image==0.16.2 scikit-learn==0.20.0 scipy==1.3.3 six==1.13.0 sklearn==0.0 threadpoolctl==1.1.0 -torch==0.4.1 +torch==1.3.1 +torchvision==0.4.2 +tqdm==4.40.2 traceback2==1.4.0 +traitlets==4.3.3 unittest2==1.1.0 +wcwidth==0.1.7 diff --git a/setup.py b/setup.py index 32dc403..8fc6cb9 100644 --- a/setup.py +++ b/setup.py @@ -1,33 +1,61 @@ #!/usr/bin/env python -# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ). +# Copyright 2019 Division of Medical Image Computing, German Cancer Research Center (DKFZ). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -from distutils.core import setup -from setuptools import find_packages +from setuptools import find_packages, setup +import os + +def parse_requirements(filename, exclude=[]): + lineiter = (line.strip() for line in open(filename)) + return [line for line in lineiter if line and not line.startswith("#") and not line.split("==")[0] in exclude] + +def install_custom_ext(setup_path): + os.system("python "+setup_path+" install") + return + +def clean(): + """Custom clean command to tidy up the project root.""" + os.system('rm -vrf ./build ./dist ./*.pyc ./*.tgz ./*.egg-info') req_file = "requirements.txt" +custom_exts = ["nms-extension", "RoIAlign-extension-2D", "RoIAlign-extension-3D"] +install_reqs = parse_requirements(req_file, exclude=custom_exts) -def parse_requirements(filename): - lineiter = (line.strip() for line in open(filename)) - return [line for line in lineiter if line and not line.startswith("#")] -install_reqs = parse_requirements(req_file) setup(name='medicaldetectiontoolkit', version='0.0.1', + url="https://github.com/MIC-DKFZ/medicaldetectiontoolkit", + author='P. Jaeger, G. Ramien, MIC at DKFZ Heidelberg', + licence="Apache 2.0", + description="Medical Object-Detection Toolkit.", + classifiers=[ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3.7" + ], packages=find_packages(exclude=['test', 'test.*']), install_requires=install_reqs, - dependency_links=[], - ) \ No newline at end of file + ) + +custom_exts = ["custom_extensions/nms", "custom_extensions/roi_align"] +for path in custom_exts: + setup_path = os.path.join(path, "setup.py") + try: + install_custom_ext(setup_path) + except Exception as e: + print("FAILED to install custom extension {} due to Error:\n{}".format(path, e)) + +clean() \ No newline at end of file diff --git a/unittests.py b/unittests.py new file mode 100644 index 0000000..41dab33 --- /dev/null +++ b/unittests.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python +# Copyright 2019 Division of Medical Image Computing, German Cancer Research Center (DKFZ). +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import unittest + +import os +import pickle +import time +from multiprocessing import Pool +import subprocess + +import numpy as np +import pandas as pd +import torch +import torchvision as tv + +import tqdm + +import utils.exp_utils as utils +import utils.model_utils as mutils + +""" Note on unittests: run this file either in the way intended for unittests by starting the script with + python -m unittest unittests.py or start it as a normal python file as python unittests.py. + You can selective run single tests by calling python -m unittest unittests.TestClassOfYourChoice, where + TestClassOfYourChoice is the name of the test defined below, e.g., CompareFoldSplits. +""" + + + +def inspect_info_df(pp_dir): + """ use your debugger to look into the info df of a pp dir. + :param pp_dir: preprocessed-data directory + """ + + info_df = pd.read_pickle(os.path.join(pp_dir, "info_df.pickle")) + + return + + +def generate_boxes(count, dim=2, h=100, w=100, d=20, normalize=False, on_grid=False, seed=0): + """ generate boxes of format [y1, x1, y2, x2, (z1, z2)]. + :param count: nr of boxes + :param dim: dimension of boxes (2 or 3) + :return: boxes in format (n_boxes, 4 or 6), scores + """ + np.random.seed(seed) + if on_grid: + lower_y = np.random.randint(0, h // 2, (count,)) + lower_x = np.random.randint(0, w // 2, (count,)) + upper_y = np.random.randint(h // 2, h, (count,)) + upper_x = np.random.randint(w // 2, w, (count,)) + if dim == 3: + lower_z = np.random.randint(0, d // 2, (count,)) + upper_z = np.random.randint(d // 2, d, (count,)) + else: + lower_y = np.random.rand(count) * h / 2. + lower_x = np.random.rand(count) * w / 2. + upper_y = (np.random.rand(count) + 1.) * h / 2. + upper_x = (np.random.rand(count) + 1.) * w / 2. + if dim == 3: + lower_z = np.random.rand(count) * d / 2. + upper_z = (np.random.rand(count) + 1.) * d / 2. + + if dim == 3: + boxes = np.array(list(zip(lower_y, lower_x, upper_y, upper_x, lower_z, upper_z))) + # add an extreme box that tests the boundaries + boxes = np.concatenate((boxes, np.array([[0., 0., h, w, 0, d]]))) + else: + boxes = np.array(list(zip(lower_y, lower_x, upper_y, upper_x))) + boxes = np.concatenate((boxes, np.array([[0., 0., h, w]]))) + + scores = np.random.rand(count + 1) + if normalize: + divisor = np.array([h, w, h, w, d, d]) if dim == 3 else np.array([h, w, h, w]) + boxes = boxes / divisor + return boxes, scores + + +# -------- check own nms CUDA implement against own numpy implement ------ +class CheckNMSImplementation(unittest.TestCase): + + @staticmethod + def assert_res_equality(keep_ics1, keep_ics2, boxes, scores, tolerance=0, names=("res1", "res2")): + """ + :param keep_ics1: keep indices (results), torch.Tensor of shape (n_ics,) + :param keep_ics2: + :return: + """ + keep_ics1, keep_ics2 = keep_ics1.cpu().numpy(), keep_ics2.cpu().numpy() + discrepancies = np.setdiff1d(keep_ics1, keep_ics2) + try: + checks = np.array([ + len(discrepancies) <= tolerance + ]) + except: + checks = np.zeros((1,)).astype("bool") + msgs = np.array([ + """{}: {} \n{}: {} \nboxes: {}\n {}\n""".format(names[0], keep_ics1, names[1], keep_ics2, boxes, + scores) + ]) + + assert np.all(checks), "NMS: results mismatch: " + "\n".join(msgs[~checks]) + + def single_case(self, count=20, dim=3, threshold=0.2, seed=0): + boxes, scores = generate_boxes(count, dim, seed=seed, h=320, w=280, d=30) + + keep_numpy = torch.tensor(mutils.nms_numpy(boxes, scores, threshold)) + + # for some reason torchvision nms requires box coords as floats. + boxes = torch.from_numpy(boxes).type(torch.float32) + scores = torch.from_numpy(scores).type(torch.float32) + if dim == 2: + """need to wait until next pytorch release where they fixed nms on cpu (currently they have >= where it + needs to be >. + """ + # keep_ops = tv.ops.nms(boxes, scores, threshold) + # self.assert_res_equality(keep_numpy, keep_ops, boxes, scores, tolerance=0, names=["np", "ops"]) + pass + + boxes = boxes.cuda() + scores = scores.cuda() + keep = self.nms_ext.nms(boxes, scores, threshold) + self.assert_res_equality(keep_numpy, keep, boxes, scores, tolerance=0, names=["np", "cuda"]) + + def test(self, n_cases=200, box_count=30, threshold=0.5): + # dynamically import module so that it doesn't affect other tests if import fails + self.nms_ext = utils.import_module("nms_ext", 'custom_extensions/nms/nms.py') + # change seed to something fix if you want exactly reproducible test + seed0 = np.random.randint(50) + print("NMS test progress (done/total box configurations) 2D:", end="\n") + for i in tqdm.tqdm(range(n_cases)): + self.single_case(count=box_count, dim=2, threshold=threshold, seed=seed0+i) + print("NMS test progress (done/total box configurations) 3D:", end="\n") + for i in tqdm.tqdm(range(n_cases)): + self.single_case(count=box_count, dim=3, threshold=threshold, seed=seed0+i) + + return + +class CheckRoIAlignImplementation(unittest.TestCase): + + def prepare(self, dim=2): + + b, c, h, w = 1, 3, 50, 50 + # feature map, (b, c, h, w(, z)) + if dim == 2: + fmap = torch.rand(b, c, h, w).cuda() + # rois = torch.tensor([[ + # [0.1, 0.1, 0.3, 0.3], + # [0.2, 0.2, 0.4, 0.7], + # [0.5, 0.7, 0.7, 0.9], + # ]]).cuda() + pool_size = (7, 7) + rois = generate_boxes(5, dim=dim, h=h, w=w, on_grid=True, seed=np.random.randint(50))[0] + elif dim == 3: + d = 20 + fmap = torch.rand(b, c, h, w, d).cuda() + # rois = torch.tensor([[ + # [0.1, 0.1, 0.3, 0.3, 0.1, 0.1], + # [0.2, 0.2, 0.4, 0.7, 0.2, 0.4], + # [0.5, 0.0, 0.7, 1.0, 0.4, 0.5], + # [0.0, 0.0, 0.9, 1.0, 0.0, 1.0], + # ]]).cuda() + pool_size = (7, 7, 3) + rois = generate_boxes(5, dim=dim, h=h, w=w, d=d, on_grid=True, seed=np.random.randint(50), + normalize=False)[0] + else: + raise ValueError("dim needs to be 2 or 3") + + rois = [torch.from_numpy(rois).type(dtype=torch.float32).cuda(), ] + fmap.requires_grad_(True) + return fmap, rois, pool_size + + def check_2d(self): + + fmap, rois, pool_size = self.prepare(dim=2) + align_ops = tv.ops.roi_align(fmap, rois, pool_size) + loss_ops = align_ops.sum() + loss_ops.backward() + + ra_object = self.ra_ext.RoIAlign(output_size=pool_size, spatial_scale=1., sampling_ratio=-1) + align_ext = ra_object(fmap, rois) + loss_ext = align_ext.sum() + loss_ext.backward() + assert (loss_ops == loss_ext), "sum of roialign ops and extension 2D diverges" + assert (align_ops == align_ext).all(), "ROIAlign failed 2D test" + + def check_3d(self): + fmap, rois, pool_size = self.prepare(dim=3) + ra_object = self.ra_ext.RoIAlign(output_size=pool_size, spatial_scale=1., sampling_ratio=-1) + align_ext = ra_object(fmap, rois) + loss_ext = align_ext.sum() + loss_ext.backward() + + align_np = mutils.roi_align_3d_numpy(fmap.cpu().detach().numpy(), [roi.cpu().numpy() for roi in rois], + pool_size) + align_np = np.squeeze(align_np) # remove singleton batch dim + + align_ext = align_ext.cpu().detach().numpy() + assert np.allclose(align_np, align_ext, rtol=1e-5, + atol=1e-8), "RoIAlign differences in numpy and CUDA implement" + + def test(self): + # dynamically import module so that it doesn't affect other tests if import fails + self.ra_ext = utils.import_module("ra_ext", 'custom_extensions/roi_align/roi_align.py') + + # 2d test + self.check_2d() + + # 3d test + self.check_3d() + + return + + +if __name__=="__main__": + stime = time.time() + + unittest.main() + + mins, secs = divmod((time.time() - stime), 60) + h, mins = divmod(mins, 60) + t = "{:d}h:{:02d}m:{:02d}s".format(int(h), int(mins), int(secs)) + print("{} total runtime: {}".format(os.path.split(__file__)[1], t)) \ No newline at end of file diff --git a/utils/model_utils.py b/utils/model_utils.py index dfd40d3..3251577 100644 --- a/utils/model_utils.py +++ b/utils/model_utils.py @@ -1,891 +1,1011 @@ #!/usr/bin/env python # Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """ Parts are based on https://github.com/multimodallearning/pytorch-mask-rcnn published under MIT license. """ import numpy as np import scipy.misc import scipy.ndimage +import scipy.interpolate import torch from torch.autograd import Variable import torch.nn as nn - +import tqdm ############################################################ # Bounding Boxes ############################################################ def compute_iou_2D(box, boxes, box_area, boxes_area): """Calculates IoU of the given box with the array of the given boxes. box: 1D vector [y1, x1, y2, x2] THIS IS THE GT BOX boxes: [boxes_count, (y1, x1, y2, x2)] box_area: float. the area of 'box' boxes_area: array of length boxes_count. Note: the areas are passed in rather than calculated here for efficency. Calculate once in the caller to avoid duplicate work. """ # Calculate intersection areas y1 = np.maximum(box[0], boxes[:, 0]) y2 = np.minimum(box[2], boxes[:, 2]) x1 = np.maximum(box[1], boxes[:, 1]) x2 = np.minimum(box[3], boxes[:, 3]) intersection = np.maximum(x2 - x1, 0) * np.maximum(y2 - y1, 0) union = box_area + boxes_area[:] - intersection[:] iou = intersection / union return iou def compute_iou_3D(box, boxes, box_volume, boxes_volume): """Calculates IoU of the given box with the array of the given boxes. box: 1D vector [y1, x1, y2, x2, z1, z2] (typically gt box) boxes: [boxes_count, (y1, x1, y2, x2, z1, z2)] box_area: float. the area of 'box' boxes_area: array of length boxes_count. Note: the areas are passed in rather than calculated here for efficency. Calculate once in the caller to avoid duplicate work. """ # Calculate intersection areas y1 = np.maximum(box[0], boxes[:, 0]) y2 = np.minimum(box[2], boxes[:, 2]) x1 = np.maximum(box[1], boxes[:, 1]) x2 = np.minimum(box[3], boxes[:, 3]) z1 = np.maximum(box[4], boxes[:, 4]) z2 = np.minimum(box[5], boxes[:, 5]) intersection = np.maximum(x2 - x1, 0) * np.maximum(y2 - y1, 0) * np.maximum(z2 - z1, 0) union = box_volume + boxes_volume[:] - intersection[:] iou = intersection / union return iou def compute_overlaps(boxes1, boxes2): """Computes IoU overlaps between two sets of boxes. boxes1, boxes2: [N, (y1, x1, y2, x2)]. / 3D: (z1, z2)) For better performance, pass the largest set first and the smaller second. """ # Areas of anchors and GT boxes if boxes1.shape[1] == 4: area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1]) area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1]) # Compute overlaps to generate matrix [boxes1 count, boxes2 count] # Each cell contains the IoU value. overlaps = np.zeros((boxes1.shape[0], boxes2.shape[0])) for i in range(overlaps.shape[1]): box2 = boxes2[i] #this is the gt box overlaps[:, i] = compute_iou_2D(box2, boxes1, area2[i], area1) return overlaps else: # Areas of anchors and GT boxes volume1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1]) * (boxes1[:, 5] - boxes1[:, 4]) volume2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1]) * (boxes2[:, 5] - boxes2[:, 4]) # Compute overlaps to generate matrix [boxes1 count, boxes2 count] # Each cell contains the IoU value. overlaps = np.zeros((boxes1.shape[0], boxes2.shape[0])) for i in range(overlaps.shape[1]): box2 = boxes2[i] # this is the gt box overlaps[:, i] = compute_iou_3D(box2, boxes1, volume2[i], volume1) return overlaps def box_refinement(box, gt_box): """Compute refinement needed to transform box to gt_box. box and gt_box are [N, (y1, x1, y2, x2)] / 3D: (z1, z2)) """ height = box[:, 2] - box[:, 0] width = box[:, 3] - box[:, 1] center_y = box[:, 0] + 0.5 * height center_x = box[:, 1] + 0.5 * width gt_height = gt_box[:, 2] - gt_box[:, 0] gt_width = gt_box[:, 3] - gt_box[:, 1] gt_center_y = gt_box[:, 0] + 0.5 * gt_height gt_center_x = gt_box[:, 1] + 0.5 * gt_width dy = (gt_center_y - center_y) / height dx = (gt_center_x - center_x) / width dh = torch.log(gt_height / height) dw = torch.log(gt_width / width) result = torch.stack([dy, dx, dh, dw], dim=1) if box.shape[1] > 4: depth = box[:, 5] - box[:, 4] center_z = box[:, 4] + 0.5 * depth gt_depth = gt_box[:, 5] - gt_box[:, 4] gt_center_z = gt_box[:, 4] + 0.5 * gt_depth dz = (gt_center_z - center_z) / depth dd = torch.log(gt_depth / depth) result = torch.stack([dy, dx, dz, dh, dw, dd], dim=1) return result def unmold_mask_2D(mask, bbox, image_shape): """Converts a mask generated by the neural network into a format similar to it's original shape. mask: [height, width] of type float. A small, typically 28x28 mask. bbox: [y1, x1, y2, x2]. The box to fit the mask in. Returns a binary mask with the same size as the original image. """ y1, x1, y2, x2 = bbox out_zoom = [y2 - y1, x2 - x1] zoom_factor = [i / j for i, j in zip(out_zoom, mask.shape)] mask = scipy.ndimage.zoom(mask, zoom_factor, order=1).astype(np.float32) # Put the mask in the right location. full_mask = np.zeros(image_shape[:2]) full_mask[y1:y2, x1:x2] = mask return full_mask def unmold_mask_3D(mask, bbox, image_shape): """Converts a mask generated by the neural network into a format similar to it's original shape. mask: [height, width] of type float. A small, typically 28x28 mask. bbox: [y1, x1, y2, x2, z1, z2]. The box to fit the mask in. Returns a binary mask with the same size as the original image. """ y1, x1, y2, x2, z1, z2 = bbox out_zoom = [y2 - y1, x2 - x1, z2 - z1] zoom_factor = [i/j for i,j in zip(out_zoom, mask.shape)] mask = scipy.ndimage.zoom(mask, zoom_factor, order=1).astype(np.float32) # Put the mask in the right location. full_mask = np.zeros(image_shape[:3]) full_mask[y1:y2, x1:x2, z1:z2] = mask return full_mask ############################################################ # Anchors ############################################################ def generate_anchors(scales, ratios, shape, feature_stride, anchor_stride): """ scales: 1D array of anchor sizes in pixels. Example: [32, 64, 128] ratios: 1D array of anchor ratios of width/height. Example: [0.5, 1, 2] shape: [height, width] spatial shape of the feature map over which to generate anchors. feature_stride: Stride of the feature map relative to the image in pixels. anchor_stride: Stride of anchors on the feature map. For example, if the value is 2 then generate anchors for every other feature map pixel. """ # Get all combinations of scales and ratios scales, ratios = np.meshgrid(np.array(scales), np.array(ratios)) scales = scales.flatten() ratios = ratios.flatten() # Enumerate heights and widths from scales and ratios heights = scales / np.sqrt(ratios) widths = scales * np.sqrt(ratios) # Enumerate shifts in feature space shifts_y = np.arange(0, shape[0], anchor_stride) * feature_stride shifts_x = np.arange(0, shape[1], anchor_stride) * feature_stride shifts_x, shifts_y = np.meshgrid(shifts_x, shifts_y) # Enumerate combinations of shifts, widths, and heights box_widths, box_centers_x = np.meshgrid(widths, shifts_x) box_heights, box_centers_y = np.meshgrid(heights, shifts_y) # Reshape to get a list of (y, x) and a list of (h, w) box_centers = np.stack( [box_centers_y, box_centers_x], axis=2).reshape([-1, 2]) box_sizes = np.stack([box_heights, box_widths], axis=2).reshape([-1, 2]) # Convert to corner coordinates (y1, x1, y2, x2) boxes = np.concatenate([box_centers - 0.5 * box_sizes, box_centers + 0.5 * box_sizes], axis=1) return boxes def generate_anchors_3D(scales_xy, scales_z, ratios, shape, feature_stride_xy, feature_stride_z, anchor_stride): """ scales: 1D array of anchor sizes in pixels. Example: [32, 64, 128] ratios: 1D array of anchor ratios of width/height. Example: [0.5, 1, 2] shape: [height, width] spatial shape of the feature map over which to generate anchors. feature_stride: Stride of the feature map relative to the image in pixels. anchor_stride: Stride of anchors on the feature map. For example, if the value is 2 then generate anchors for every other feature map pixel. """ # Get all combinations of scales and ratios scales_xy, ratios_meshed = np.meshgrid(np.array(scales_xy), np.array(ratios)) scales_xy = scales_xy.flatten() ratios_meshed = ratios_meshed.flatten() # Enumerate heights and widths from scales and ratios heights = scales_xy / np.sqrt(ratios_meshed) widths = scales_xy * np.sqrt(ratios_meshed) depths = np.tile(np.array(scales_z), len(ratios_meshed)//np.array(scales_z)[..., None].shape[0]) # Enumerate shifts in feature space shifts_y = np.arange(0, shape[0], anchor_stride) * feature_stride_xy #translate from fm positions to input coords. shifts_x = np.arange(0, shape[1], anchor_stride) * feature_stride_xy shifts_z = np.arange(0, shape[2], anchor_stride) * (feature_stride_z) shifts_x, shifts_y, shifts_z = np.meshgrid(shifts_x, shifts_y, shifts_z) # Enumerate combinations of shifts, widths, and heights box_widths, box_centers_x = np.meshgrid(widths, shifts_x) box_heights, box_centers_y = np.meshgrid(heights, shifts_y) box_depths, box_centers_z = np.meshgrid(depths, shifts_z) # Reshape to get a list of (y, x, z) and a list of (h, w, d) box_centers = np.stack( [box_centers_y, box_centers_x, box_centers_z], axis=2).reshape([-1, 3]) box_sizes = np.stack([box_heights, box_widths, box_depths], axis=2).reshape([-1, 3]) # Convert to corner coordinates (y1, x1, y2, x2, z1, z2) boxes = np.concatenate([box_centers - 0.5 * box_sizes, box_centers + 0.5 * box_sizes], axis=1) boxes = np.transpose(np.array([boxes[:, 0], boxes[:, 1], boxes[:, 3], boxes[:, 4], boxes[:, 2], boxes[:, 5]]), axes=(1, 0)) return boxes def generate_pyramid_anchors(logger, cf): """Generate anchors at different levels of a feature pyramid. Each scale is associated with a level of the pyramid, but each ratio is used in all levels of the pyramid. from configs: :param scales: cf.RPN_ANCHOR_SCALES , e.g. [4, 8, 16, 32] :param ratios: cf.RPN_ANCHOR_RATIOS , e.g. [0.5, 1, 2] :param feature_shapes: cf.BACKBONE_SHAPES , e.g. [array of shapes per feature map] [80, 40, 20, 10, 5] :param feature_strides: cf.BACKBONE_STRIDES , e.g. [2, 4, 8, 16, 32, 64] :param anchors_stride: cf.RPN_ANCHOR_STRIDE , e.g. 1 :return anchors: (N, (y1, x1, y2, x2, (z1), (z2)). All generated anchors in one array. Sorted with the same order of the given scales. So, anchors of scale[0] come first, then anchors of scale[1], and so on. """ scales = cf.rpn_anchor_scales ratios = cf.rpn_anchor_ratios feature_shapes = cf.backbone_shapes anchor_stride = cf.rpn_anchor_stride pyramid_levels = cf.pyramid_levels feature_strides = cf.backbone_strides anchors = [] logger.info("feature map shapes: {}".format(feature_shapes)) logger.info("anchor scales: {}".format(scales)) expected_anchors = [np.prod(feature_shapes[ii]) * len(ratios) * len(scales['xy'][ii]) for ii in pyramid_levels] for lix, level in enumerate(pyramid_levels): if len(feature_shapes[level]) == 2: anchors.append(generate_anchors(scales['xy'][level], ratios, feature_shapes[level], feature_strides['xy'][level], anchor_stride)) else: anchors.append(generate_anchors_3D(scales['xy'][level], scales['z'][level], ratios, feature_shapes[level], feature_strides['xy'][level], feature_strides['z'][level], anchor_stride)) logger.info("level {}: built anchors {} / expected anchors {} ||| total build {} / total expected {}".format( level, anchors[-1].shape, expected_anchors[lix], np.concatenate(anchors).shape, np.sum(expected_anchors))) out_anchors = np.concatenate(anchors, axis=0) return out_anchors def apply_box_deltas_2D(boxes, deltas): """Applies the given deltas to the given boxes. boxes: [N, 4] where each row is y1, x1, y2, x2 deltas: [N, 4] where each row is [dy, dx, log(dh), log(dw)] """ # Convert to y, x, h, w height = boxes[:, 2] - boxes[:, 0] width = boxes[:, 3] - boxes[:, 1] center_y = boxes[:, 0] + 0.5 * height center_x = boxes[:, 1] + 0.5 * width # Apply deltas center_y += deltas[:, 0] * height center_x += deltas[:, 1] * width height *= torch.exp(deltas[:, 2]) width *= torch.exp(deltas[:, 3]) # Convert back to y1, x1, y2, x2 y1 = center_y - 0.5 * height x1 = center_x - 0.5 * width y2 = y1 + height x2 = x1 + width result = torch.stack([y1, x1, y2, x2], dim=1) return result def apply_box_deltas_3D(boxes, deltas): """Applies the given deltas to the given boxes. boxes: [N, 6] where each row is y1, x1, y2, x2, z1, z2 deltas: [N, 6] where each row is [dy, dx, dz, log(dh), log(dw), log(dd)] """ # Convert to y, x, h, w height = boxes[:, 2] - boxes[:, 0] width = boxes[:, 3] - boxes[:, 1] depth = boxes[:, 5] - boxes[:, 4] center_y = boxes[:, 0] + 0.5 * height center_x = boxes[:, 1] + 0.5 * width center_z = boxes[:, 4] + 0.5 * depth # Apply deltas center_y += deltas[:, 0] * height center_x += deltas[:, 1] * width center_z += deltas[:, 2] * depth height *= torch.exp(deltas[:, 3]) width *= torch.exp(deltas[:, 4]) depth *= torch.exp(deltas[:, 5]) # Convert back to y1, x1, y2, x2 y1 = center_y - 0.5 * height x1 = center_x - 0.5 * width z1 = center_z - 0.5 * depth y2 = y1 + height x2 = x1 + width z2 = z1 + depth result = torch.stack([y1, x1, y2, x2, z1, z2], dim=1) return result def clip_boxes_2D(boxes, window): """ boxes: [N, 4] each col is y1, x1, y2, x2 window: [4] in the form y1, x1, y2, x2 """ boxes = torch.stack( \ [boxes[:, 0].clamp(float(window[0]), float(window[2])), boxes[:, 1].clamp(float(window[1]), float(window[3])), boxes[:, 2].clamp(float(window[0]), float(window[2])), boxes[:, 3].clamp(float(window[1]), float(window[3]))], 1) return boxes def clip_boxes_3D(boxes, window): """ boxes: [N, 6] each col is y1, x1, y2, x2, z1, z2 window: [6] in the form y1, x1, y2, x2, z1, z2 """ boxes = torch.stack( \ [boxes[:, 0].clamp(float(window[0]), float(window[2])), boxes[:, 1].clamp(float(window[1]), float(window[3])), boxes[:, 2].clamp(float(window[0]), float(window[2])), boxes[:, 3].clamp(float(window[1]), float(window[3])), boxes[:, 4].clamp(float(window[4]), float(window[5])), boxes[:, 5].clamp(float(window[4]), float(window[5]))], 1) return boxes def clip_boxes_numpy(boxes, window): """ boxes: [N, 4] each col is y1, x1, y2, x2 / [N, 6] in 3D. window: iamge shape (y, x, (z)) """ if boxes.shape[1] == 4: boxes = np.concatenate( (np.clip(boxes[:, 0], 0, window[0])[:, None], np.clip(boxes[:, 1], 0, window[0])[:, None], np.clip(boxes[:, 2], 0, window[1])[:, None], np.clip(boxes[:, 3], 0, window[1])[:, None]), 1 ) else: boxes = np.concatenate( (np.clip(boxes[:, 0], 0, window[0])[:, None], np.clip(boxes[:, 1], 0, window[0])[:, None], np.clip(boxes[:, 2], 0, window[1])[:, None], np.clip(boxes[:, 3], 0, window[1])[:, None], np.clip(boxes[:, 4], 0, window[2])[:, None], np.clip(boxes[:, 5], 0, window[2])[:, None]), 1 ) return boxes def bbox_overlaps_2D(boxes1, boxes2): """Computes IoU overlaps between two sets of boxes. boxes1, boxes2: [N, (y1, x1, y2, x2)]. """ # 1. Tile boxes2 and repeate boxes1. This allows us to compare # every boxes1 against every boxes2 without loops. # TF doesn't have an equivalent to np.repeate() so simulate it # using tf.tile() and tf.reshape. boxes1_repeat = boxes2.size()[0] boxes2_repeat = boxes1.size()[0] boxes1 = boxes1.repeat(1,boxes1_repeat).view(-1,4) boxes2 = boxes2.repeat(boxes2_repeat,1) # 2. Compute intersections b1_y1, b1_x1, b1_y2, b1_x2 = boxes1.chunk(4, dim=1) b2_y1, b2_x1, b2_y2, b2_x2 = boxes2.chunk(4, dim=1) y1 = torch.max(b1_y1, b2_y1)[:, 0] x1 = torch.max(b1_x1, b2_x1)[:, 0] y2 = torch.min(b1_y2, b2_y2)[:, 0] x2 = torch.min(b1_x2, b2_x2)[:, 0] zeros = Variable(torch.zeros(y1.size()[0]), requires_grad=False) if y1.is_cuda: zeros = zeros.cuda() intersection = torch.max(x2 - x1, zeros) * torch.max(y2 - y1, zeros) # 3. Compute unions b1_area = (b1_y2 - b1_y1) * (b1_x2 - b1_x1) b2_area = (b2_y2 - b2_y1) * (b2_x2 - b2_x1) union = b1_area[:,0] + b2_area[:,0] - intersection # 4. Compute IoU and reshape to [boxes1, boxes2] iou = intersection / union overlaps = iou.view(boxes2_repeat, boxes1_repeat) return overlaps def bbox_overlaps_3D(boxes1, boxes2): """Computes IoU overlaps between two sets of boxes. boxes1, boxes2: [N, (y1, x1, y2, x2, z1, z2)]. """ # 1. Tile boxes2 and repeate boxes1. This allows us to compare # every boxes1 against every boxes2 without loops. # TF doesn't have an equivalent to np.repeate() so simulate it # using tf.tile() and tf.reshape. boxes1_repeat = boxes2.size()[0] boxes2_repeat = boxes1.size()[0] boxes1 = boxes1.repeat(1,boxes1_repeat).view(-1,6) boxes2 = boxes2.repeat(boxes2_repeat,1) # 2. Compute intersections b1_y1, b1_x1, b1_y2, b1_x2, b1_z1, b1_z2 = boxes1.chunk(6, dim=1) b2_y1, b2_x1, b2_y2, b2_x2, b2_z1, b2_z2 = boxes2.chunk(6, dim=1) y1 = torch.max(b1_y1, b2_y1)[:, 0] x1 = torch.max(b1_x1, b2_x1)[:, 0] y2 = torch.min(b1_y2, b2_y2)[:, 0] x2 = torch.min(b1_x2, b2_x2)[:, 0] z1 = torch.max(b1_z1, b2_z1)[:, 0] z2 = torch.min(b1_z2, b2_z2)[:, 0] zeros = Variable(torch.zeros(y1.size()[0]), requires_grad=False) if y1.is_cuda: zeros = zeros.cuda() intersection = torch.max(x2 - x1, zeros) * torch.max(y2 - y1, zeros) * torch.max(z2 - z1, zeros) # 3. Compute unions b1_volume = (b1_y2 - b1_y1) * (b1_x2 - b1_x1) * (b1_z2 - b1_z1) b2_volume = (b2_y2 - b2_y1) * (b2_x2 - b2_x1) * (b2_z2 - b2_z1) union = b1_volume[:,0] + b2_volume[:,0] - intersection # 4. Compute IoU and reshape to [boxes1, boxes2] iou = intersection / union overlaps = iou.view(boxes2_repeat, boxes1_repeat) return overlaps def gt_anchor_matching(cf, anchors, gt_boxes, gt_class_ids=None): """Given the anchors and GT boxes, compute overlaps and identify positive anchors and deltas to refine them to match their corresponding GT boxes. anchors: [num_anchors, (y1, x1, y2, x2, (z1), (z2))] gt_boxes: [num_gt_boxes, (y1, x1, y2, x2, (z1), (z2))] gt_class_ids (optional): [num_gt_boxes] Integer class IDs for one stage detectors. in RPN case of Mask R-CNN, set all positive matches to 1 (foreground) Returns: anchor_class_matches: [N] (int32) matches between anchors and GT boxes. 1 = positive anchor, -1 = negative anchor, 0 = neutral. In case of one stage detectors like RetinaNet/RetinaUNet this flag takes class_ids as positive anchor values, i.e. values >= 1! anchor_delta_targets: [N, (dy, dx, (dz), log(dh), log(dw), (log(dd)))] Anchor bbox deltas. """ anchor_class_matches = np.zeros([anchors.shape[0]], dtype=np.int32) anchor_delta_targets = np.zeros((cf.rpn_train_anchors_per_image, 2*cf.dim)) anchor_matching_iou = cf.anchor_matching_iou if gt_boxes is None: anchor_class_matches = np.full(anchor_class_matches.shape, fill_value=-1) return anchor_class_matches, anchor_delta_targets # for mrcnn: anchor matching is done for RPN loss, so positive labels are all 1 (foreground) if gt_class_ids is None: gt_class_ids = np.array([1] * len(gt_boxes)) # Compute overlaps [num_anchors, num_gt_boxes] overlaps = compute_overlaps(anchors, gt_boxes) # Match anchors to GT Boxes # If an anchor overlaps a GT box with IoU >= anchor_matching_iou then it's positive. # If an anchor overlaps a GT box with IoU < 0.1 then it's negative. # Neutral anchors are those that don't match the conditions above, # and they don't influence the loss function. # However, don't keep any GT box unmatched (rare, but happens). Instead, # match it to the closest anchor (even if its max IoU is < 0.1). # 1. Set negative anchors first. They get overwritten below if a GT box is # matched to them. Skip boxes in crowd areas. anchor_iou_argmax = np.argmax(overlaps, axis=1) anchor_iou_max = overlaps[np.arange(overlaps.shape[0]), anchor_iou_argmax] if anchors.shape[1] == 4: anchor_class_matches[(anchor_iou_max < 0.1)] = -1 elif anchors.shape[1] == 6: anchor_class_matches[(anchor_iou_max < 0.01)] = -1 else: raise ValueError('anchor shape wrong {}'.format(anchors.shape)) # 2. Set an anchor for each GT box (regardless of IoU value). gt_iou_argmax = np.argmax(overlaps, axis=0) for ix, ii in enumerate(gt_iou_argmax): anchor_class_matches[ii] = gt_class_ids[ix] # 3. Set anchors with high overlap as positive. above_trhesh_ixs = np.argwhere(anchor_iou_max >= anchor_matching_iou) anchor_class_matches[above_trhesh_ixs] = gt_class_ids[anchor_iou_argmax[above_trhesh_ixs]] # Subsample to balance positive anchors. ids = np.where(anchor_class_matches > 0)[0] extra = len(ids) - (cf.rpn_train_anchors_per_image // 2) if extra > 0: # Reset the extra ones to neutral ids = np.random.choice(ids, extra, replace=False) anchor_class_matches[ids] = 0 # Leave all negative proposals negative now and sample from them in online hard example mining. # For positive anchors, compute shift and scale needed to transform them to match the corresponding GT boxes. ids = np.where(anchor_class_matches > 0)[0] ix = 0 # index into anchor_delta_targets for i, a in zip(ids, anchors[ids]): # closest gt box (it might have IoU < anchor_matching_iou) gt = gt_boxes[anchor_iou_argmax[i]] # convert coordinates to center plus width/height. gt_h = gt[2] - gt[0] gt_w = gt[3] - gt[1] gt_center_y = gt[0] + 0.5 * gt_h gt_center_x = gt[1] + 0.5 * gt_w # Anchor a_h = a[2] - a[0] a_w = a[3] - a[1] a_center_y = a[0] + 0.5 * a_h a_center_x = a[1] + 0.5 * a_w if cf.dim == 2: anchor_delta_targets[ix] = [ (gt_center_y - a_center_y) / a_h, (gt_center_x - a_center_x) / a_w, np.log(gt_h / a_h), np.log(gt_w / a_w), ] else: gt_d = gt[5] - gt[4] gt_center_z = gt[4] + 0.5 * gt_d a_d = a[5] - a[4] a_center_z = a[4] + 0.5 * a_d anchor_delta_targets[ix] = [ (gt_center_y - a_center_y) / a_h, (gt_center_x - a_center_x) / a_w, (gt_center_z - a_center_z) / a_d, np.log(gt_h / a_h), np.log(gt_w / a_w), np.log(gt_d / a_d) ] # normalize. anchor_delta_targets[ix] /= cf.rpn_bbox_std_dev ix += 1 return anchor_class_matches, anchor_delta_targets def clip_to_window(window, boxes): """ window: (y1, x1, y2, x2) / 3D: (z1, z2). The window in the image we want to clip to. boxes: [N, (y1, x1, y2, x2)] / 3D: (z1, z2) """ boxes[:, 0] = boxes[:, 0].clamp(float(window[0]), float(window[2])) boxes[:, 1] = boxes[:, 1].clamp(float(window[1]), float(window[3])) boxes[:, 2] = boxes[:, 2].clamp(float(window[0]), float(window[2])) boxes[:, 3] = boxes[:, 3].clamp(float(window[1]), float(window[3])) if boxes.shape[1] > 5: boxes[:, 4] = boxes[:, 4].clamp(float(window[4]), float(window[5])) boxes[:, 5] = boxes[:, 5].clamp(float(window[4]), float(window[5])) return boxes +def nms_numpy(box_coords, scores, thresh): + """ non-maximum suppression on 2D or 3D boxes in numpy. + :param box_coords: [y1,x1,y2,x2 (,z1,z2)] with y1<=y2, x1<=x2, z1<=z2. + :param scores: ranking scores (higher score == higher rank) of boxes. + :param thresh: IoU threshold for clustering. + :return: + """ + y1 = box_coords[:, 0] + x1 = box_coords[:, 1] + y2 = box_coords[:, 2] + x2 = box_coords[:, 3] + assert np.all(y1 <= y2) and np.all(x1 <= x2), """"the definition of the coordinates is crucially important here: + coordinates of which maxima are taken need to be the lower coordinates""" + areas = (x2 - x1) * (y2 - y1) + + is_3d = box_coords.shape[1] == 6 + if is_3d: # 3-dim case + z1 = box_coords[:, 4] + z2 = box_coords[:, 5] + assert np.all(z1<=z2), """"the definition of the coordinates is crucially important here: + coordinates of which maxima are taken need to be the lower coordinates""" + areas *= (z2 - z1) + + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: # order is the sorted index. maps order to index: order[1] = 24 means (rank1, ix 24) + i = order[0] # highest scoring element + yy1 = np.maximum(y1[i], y1[order]) # highest scoring element still in >order<, is compared to itself, that is okay. + xx1 = np.maximum(x1[i], x1[order]) + yy2 = np.minimum(y2[i], y2[order]) + xx2 = np.minimum(x2[i], x2[order]) + + h = np.maximum(0.0, yy2 - yy1) + w = np.maximum(0.0, xx2 - xx1) + inter = h * w + + if is_3d: + zz1 = np.maximum(z1[i], z1[order]) + zz2 = np.minimum(z2[i], z2[order]) + d = np.maximum(0.0, zz2 - zz1) + inter *= d + + iou = inter / (areas[i] + areas[order] - inter) + + non_matches = np.nonzero(iou <= thresh)[0] # get all elements that were not matched and discard all others. + order = order[non_matches] + keep.append(i) + + return keep + +def roi_align_3d_numpy(input: np.ndarray, rois, output_size: tuple, + spatial_scale: float = 1., sampling_ratio: int = -1) -> np.ndarray: + """ This fct mainly serves as a verification method for 3D CUDA implementation of RoIAlign, it's highly + inefficient due to the nested loops. + :param input: (ndarray[N, C, H, W, D]): input feature map + :param rois: list (N,K(n), 6), K(n) = nr of rois in batch-element n, single roi of format (y1,x1,y2,x2,z1,z2) + :param output_size: + :param spatial_scale: + :param sampling_ratio: + :return: (List[N, K(n), C, output_size[0], output_size[1], output_size[2]]) + """ + + out_height, out_width, out_depth = output_size + + coord_grid = tuple([np.linspace(0, input.shape[dim] - 1, num=input.shape[dim]) for dim in range(2, 5)]) + pooled_rois = [[]] * len(rois) + assert len(rois) == input.shape[0], "batch dim mismatch, rois: {}, input: {}".format(len(rois), input.shape[0]) + print("Numpy 3D RoIAlign progress:", end="\n") + for b in range(input.shape[0]): + for roi in tqdm.tqdm(rois[b]): + y1, x1, y2, x2, z1, z2 = np.array(roi) * spatial_scale + roi_height = max(float(y2 - y1), 1.) + roi_width = max(float(x2 - x1), 1.) + roi_depth = max(float(z2 - z1), 1.) + + if sampling_ratio <= 0: + sampling_ratio_h = int(np.ceil(roi_height / out_height)) + sampling_ratio_w = int(np.ceil(roi_width / out_width)) + sampling_ratio_d = int(np.ceil(roi_depth / out_depth)) + else: + sampling_ratio_h = sampling_ratio_w = sampling_ratio_d = sampling_ratio # == n points per bin + + bin_height = roi_height / out_height + bin_width = roi_width / out_width + bin_depth = roi_depth / out_depth + + n_points = sampling_ratio_h * sampling_ratio_w * sampling_ratio_d + pooled_roi = np.empty((input.shape[1], out_height, out_width, out_depth), dtype="float32") + for chan in range(input.shape[1]): + lin_interpolator = scipy.interpolate.RegularGridInterpolator(coord_grid, input[b, chan], + method="linear") + for bin_iy in range(out_height): + for bin_ix in range(out_width): + for bin_iz in range(out_depth): + + bin_val = 0. + for i in range(sampling_ratio_h): + for j in range(sampling_ratio_w): + for k in range(sampling_ratio_d): + loc_ijk = [ + y1 + bin_iy * bin_height + (i + 0.5) * (bin_height / sampling_ratio_h), + x1 + bin_ix * bin_width + (j + 0.5) * (bin_width / sampling_ratio_w), + z1 + bin_iz * bin_depth + (k + 0.5) * (bin_depth / sampling_ratio_d)] + # print("loc_ijk", loc_ijk) + if not (np.any([c < -1.0 for c in loc_ijk]) or loc_ijk[0] > input.shape[2] or + loc_ijk[1] > input.shape[3] or loc_ijk[2] > input.shape[4]): + for catch_case in range(3): + # catch on-border cases + if int(loc_ijk[catch_case]) == input.shape[catch_case + 2] - 1: + loc_ijk[catch_case] = input.shape[catch_case + 2] - 1 + bin_val += lin_interpolator(loc_ijk) + pooled_roi[chan, bin_iy, bin_ix, bin_iz] = bin_val / n_points + + pooled_rois[b].append(pooled_roi) + + return np.array(pooled_rois) + + ############################################################ # Pytorch Utility Functions ############################################################ def unique1d(tensor): if tensor.size()[0] == 0 or tensor.size()[0] == 1: return tensor tensor = tensor.sort()[0] unique_bool = tensor[1:] != tensor [:-1] first_element = torch.tensor([True], dtype=torch.bool, requires_grad=False) if tensor.is_cuda: first_element = first_element.cuda() unique_bool = torch.cat((first_element, unique_bool),dim=0) return tensor[unique_bool.data] def log2(x): """Implementatin of Log2. Pytorch doesn't have a native implemenation.""" ln2 = Variable(torch.log(torch.FloatTensor([2.0])), requires_grad=False) if x.is_cuda: ln2 = ln2.cuda() return torch.log(x) / ln2 def intersect1d(tensor1, tensor2): aux = torch.cat((tensor1, tensor2), dim=0) aux = aux.sort(descending=True)[0] return aux[:-1][(aux[1:] == aux[:-1]).data] def shem(roi_probs_neg, negative_count, ohem_poolsize): """ stochastic hard example mining: from a list of indices (referring to non-matched predictions), determine a pool of highest scoring (worst false positives) of size negative_count*ohem_poolsize. Then, sample n (= negative_count) predictions of this pool as negative examples for loss. :param roi_probs_neg: tensor of shape (n_predictions, n_classes). :param negative_count: int. :param ohem_poolsize: int. :return: (negative_count). indices refer to the positions in roi_probs_neg. If pool smaller than expected due to limited negative proposals availabel, this function will return sampled indices of number < negative_count without throwing an error. """ # sort according to higehst foreground score. probs, order = roi_probs_neg[:, 1:].max(1)[0].sort(descending=True) select = torch.tensor((ohem_poolsize * int(negative_count), order.size()[0])).min().int() pool_indices = order[:select] rand_idx = torch.randperm(pool_indices.size()[0]) return pool_indices[rand_idx[:negative_count].cuda()] def initialize_weights(net): """ Initialize model weights. Current Default in Pytorch (version 0.4.1) is initialization from a uniform distriubtion. Will expectably be changed to kaiming_uniform in future versions. """ init_type = net.cf.weight_init for m in [module for module in net.modules() if type(module) in [nn.Conv2d, nn.Conv3d, nn.ConvTranspose2d, nn.ConvTranspose3d, nn.Linear]]: if init_type == 'xavier_uniform': nn.init.xavier_uniform_(m.weight.data) if m.bias is not None: m.bias.data.zero_() elif init_type == 'xavier_normal': nn.init.xavier_normal_(m.weight.data) if m.bias is not None: m.bias.data.zero_() elif init_type == "kaiming_uniform": nn.init.kaiming_uniform_(m.weight.data, mode='fan_out', nonlinearity=net.cf.relu, a=0) if m.bias is not None: fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(m.weight.data) bound = 1 / np.sqrt(fan_out) nn.init.uniform_(m.bias, -bound, bound) elif init_type == "kaiming_normal": nn.init.kaiming_normal_(m.weight.data, mode='fan_out', nonlinearity=net.cf.relu, a=0) if m.bias is not None: fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(m.weight.data) bound = 1 / np.sqrt(fan_out) nn.init.normal_(m.bias, -bound, bound) class NDConvGenerator(object): """ generic wrapper around conv-layers to avoid 2D vs. 3D distinguishing in code. """ def __init__(self, dim): self.dim = dim def __call__(self, c_in, c_out, ks, pad=0, stride=1, norm=None, relu='relu'): """ :param c_in: number of in_channels. :param c_out: number of out_channels. :param ks: kernel size. :param pad: pad size. :param stride: kernel stride. :param norm: string specifying type of feature map normalization. If None, no normalization is applied. :param relu: string specifying type of nonlinearity. If None, no nonlinearity is applied. :return: convolved feature_map. """ if self.dim == 2: conv = nn.Conv2d(c_in, c_out, kernel_size=ks, padding=pad, stride=stride) if norm is not None: if norm == 'instance_norm': norm_layer = nn.InstanceNorm2d(c_out) elif norm == 'batch_norm': norm_layer = nn.BatchNorm2d(c_out) else: raise ValueError('norm type as specified in configs is not implemented...') conv = nn.Sequential(conv, norm_layer) else: conv = nn.Conv3d(c_in, c_out, kernel_size=ks, padding=pad, stride=stride) if norm is not None: if norm == 'instance_norm': norm_layer = nn.InstanceNorm3d(c_out) elif norm == 'batch_norm': norm_layer = nn.BatchNorm3d(c_out) else: raise ValueError('norm type as specified in configs is not implemented... {}'.format(norm)) conv = nn.Sequential(conv, norm_layer) if relu is not None: if relu == 'relu': relu_layer = nn.ReLU(inplace=True) elif relu == 'leaky_relu': relu_layer = nn.LeakyReLU(inplace=True) else: raise ValueError('relu type as specified in configs is not implemented...') conv = nn.Sequential(conv, relu_layer) return conv def get_one_hot_encoding(y, n_classes): """ transform a numpy label array to a one-hot array of the same shape. :param y: array of shape (b, 1, y, x, (z)). :param n_classes: int, number of classes to unfold in one-hot encoding. :return y_ohe: array of shape (b, n_classes, y, x, (z)) """ dim = len(y.shape) - 2 if dim == 2: y_ohe = np.zeros((y.shape[0], n_classes, y.shape[2], y.shape[3])).astype('int32') if dim ==3: y_ohe = np.zeros((y.shape[0], n_classes, y.shape[2], y.shape[3], y.shape[4])).astype('int32') for cl in range(n_classes): y_ohe[:, cl][y[:, 0] == cl] = 1 return y_ohe def get_dice_per_batch_and_class(pred, y, n_classes): ''' computes dice scores per batch instance and class. :param pred: prediction array of shape (b, 1, y, x, (z)) (e.g. softmax prediction with argmax over dim 1) :param y: ground truth array of shape (b, 1, y, x, (z)) (contains int [0, ..., n_classes] :param n_classes: int :return: dice scores of shape (b, c) ''' pred = get_one_hot_encoding(pred, n_classes) y = get_one_hot_encoding(y, n_classes) axes = tuple(range(2, len(pred.shape))) intersect = np.sum(pred*y, axis=axes) denominator = np.sum(pred, axis=axes)+np.sum(y, axis=axes) + 1e-8 dice = 2.0*intersect / denominator return dice def sum_tensor(input, axes, keepdim=False): axes = np.unique(axes) if keepdim: for ax in axes: input = input.sum(ax, keepdim=True) else: for ax in sorted(axes, reverse=True): input = input.sum(int(ax)) return input def batch_dice(pred, y, false_positive_weight=1.0, smooth=1e-6): ''' compute soft dice over batch. this is a differentiable score and can be used as a loss function. only dice scores of foreground classes are returned, since training typically does not benefit from explicit background optimization. Pixels of the entire batch are considered a pseudo-volume to compute dice scores of. This way, single patches with missing foreground classes can not produce faulty gradients. :param pred: (b, c, y, x, (z)), softmax probabilities (network output). (c==classes) :param y: (b, c, y, x, (z)), one-hot-encoded segmentation mask. :param false_positive_weight: float [0,1]. For weighting of imbalanced classes, reduces the penalty for false-positive pixels. Can be beneficial sometimes in data with heavy fg/bg imbalances. :return: soft dice score (float). This function discards the background score and returns the mean of foreground scores. ''' if len(pred.size()) == 4: axes = (0, 2, 3) intersect = sum_tensor(pred * y, axes, keepdim=False) denom = sum_tensor(false_positive_weight*pred + y, axes, keepdim=False) return torch.mean(( (2 * intersect + smooth) / (denom + smooth) )[1:]) # only fg dice here. elif len(pred.size()) == 5: axes = (0, 2, 3, 4) intersect = sum_tensor(pred * y, axes, keepdim=False) denom = sum_tensor(false_positive_weight*pred + y, axes, keepdim=False) return torch.mean(( (2*intersect + smooth) / (denom + smooth) )[1:]) # only fg dice here. else: raise ValueError('wrong input dimension in dice loss') def batch_dice_mask(pred, y, mask, false_positive_weight=1.0, smooth=1e-6): ''' compute soft dice over batch. this is a diffrentiable score and can be used as a loss function. only dice scores of foreground classes are returned, since training typically does not benefit from explicit background optimization. Pixels of the entire batch are considered a pseudo-volume to compute dice scores of. This way, single patches with missing foreground classes can not produce faulty gradients. :param pred: (b, c, y, x, (z)), softmax probabilities (network output). :param y: (b, c, y, x, (z)), one hote encoded segmentation mask. :param false_positive_weight: float [0,1]. For weighting of imbalanced classes, reduces the penalty for false-positive pixels. Can be beneficial sometimes in data with heavy fg/bg imbalances. :return: soft dice score (float). This function discards the background score and returns the mean of foreground scores. ''' mask = mask.unsqueeze(1).repeat(1, 2, 1, 1) if len(pred.size()) == 4: axes = (0, 2, 3) intersect = sum_tensor(pred * y * mask, axes, keepdim=False) denom = sum_tensor(false_positive_weight*pred * mask + y * mask, axes, keepdim=False) return torch.mean(( (2*intersect + smooth) / (denom + smooth))[1:]) # only fg dice here. elif len(pred.size()) == 5: axes = (0, 2, 3, 4) intersect = sum_tensor(pred * y, axes, keepdim=False) denom = sum_tensor(false_positive_weight*pred + y, axes, keepdim=False) return torch.mean(( (2*intersect + smooth) / (denom + smooth) )[1:]) # only fg dice here. else: raise ValueError('wrong input dimension in dice loss') \ No newline at end of file