├── LICENSE ├── README.md ├── setup.py ├── tfdet ├── __init__.py ├── builder │ ├── __init__.py │ ├── api.py │ └── util.py ├── callback │ ├── __init__.py │ ├── metric.py │ ├── scheduler.py │ └── util.py ├── core │ ├── __init__.py │ ├── anchor │ │ ├── __init__.py │ │ ├── generator.py │ │ └── util.py │ ├── assign │ │ ├── __init__.py │ │ ├── atss.py │ │ ├── center_region.py │ │ ├── max_iou.py │ │ ├── ota.py │ │ ├── point.py │ │ └── sampler.py │ ├── bbox │ │ ├── __init__.py │ │ ├── coder.py │ │ ├── overlap.py │ │ └── util.py │ ├── loss │ │ ├── __init__.py │ │ ├── cross_entropy.py │ │ ├── object_detection.py │ │ ├── regression.py │ │ ├── segmentation.py │ │ └── util.py │ ├── metric │ │ ├── __init__.py │ │ ├── object_detection.py │ │ └── segmentation.py │ ├── ops │ │ ├── __init__.py │ │ ├── distance.py │ │ ├── feature_extract.py │ │ ├── initializer.py │ │ ├── nms.py │ │ ├── pooling.py │ │ └── roi_extract.py │ └── util │ │ ├── __init__.py │ │ ├── log.py │ │ ├── random.py │ │ ├── tf.py │ │ └── wrapper.py ├── dataset │ ├── __init__.py │ ├── balloon.py │ ├── coco.py │ ├── dataset.py │ ├── pascal_voc.py │ ├── pipeline │ │ ├── __init__.py │ │ ├── transform.py │ │ └── util.py │ ├── transform │ │ ├── __init__.py │ │ ├── augment.py │ │ ├── bundle.py │ │ ├── common.py │ │ └── formatting.py │ ├── util │ │ ├── __init__.py │ │ ├── file.py │ │ ├── image.py │ │ ├── numpy.py │ │ └── xml.py │ └── yolo.py ├── export │ ├── __init__.py │ ├── onnx.py │ ├── tensorrt.py │ └── tf.py ├── model │ ├── __init__.py │ ├── backbone │ │ ├── __init__.py │ │ ├── convnext.py │ │ ├── darknet.py │ │ ├── densenet.py │ │ ├── effnet.py │ │ ├── mobilenet.py │ │ ├── resnest.py │ │ ├── resnet.py │ │ ├── swin_transformer.py │ │ └── vgg.py │ ├── detector │ │ ├── __init__.py │ │ ├── deeplab.py │ │ ├── effdet.py │ │ ├── fcn.py │ │ ├── fcos.py │ │ ├── padim.py │ │ ├── patch_core.py │ │ ├── pspnet.py │ │ ├── rcnn.py │ │ ├── retina.py │ │ ├── spade.py │ │ ├── unet.py │ │ ├── upernet.py │ │ └── yolo.py │ ├── head │ │ ├── __init__.py │ │ ├── deeplab.py │ │ ├── fcn.py │ │ ├── fcos.py │ │ ├── padim.py │ │ ├── patch_core.py │ │ ├── pspnet.py │ │ ├── rcnn.py │ │ ├── retina.py │ │ ├── spade.py │ │ ├── unet.py │ │ ├── upernet.py │ │ └── yolo.py │ ├── neck │ │ ├── __init__.py │ │ ├── align.py │ │ └── fpn.py │ ├── postprocess │ │ ├── __init__.py │ │ ├── anchor.py │ │ ├── anchor_free.py │ │ ├── anodet.py │ │ ├── fcos.py │ │ ├── rcnn.py │ │ ├── retina.py │ │ ├── roi.py │ │ └── yolo.py │ └── train │ │ ├── __init__.py │ │ ├── fcos.py │ │ ├── loss │ │ ├── __init__.py │ │ ├── anchor.py │ │ ├── anchor_free.py │ │ ├── mask.py │ │ ├── roi.py │ │ ├── util.py │ │ └── yolo.py │ │ ├── padim.py │ │ ├── patch_core.py │ │ ├── rcnn.py │ │ ├── retina.py │ │ ├── segmentation.py │ │ ├── spade.py │ │ └── yolo.py └── util │ ├── __init__.py │ ├── metric.py │ └── visualize.py └── usage ├── dataset ├── custom_dataset.ipynb ├── dataset.ipynb └── yolo_dataset.ipynb ├── export ├── common.ipynb └── object_detection.ipynb ├── model ├── anomaly_detection.ipynb ├── object_detection.ipynb └── semantic_segmentation.ipynb └── train ├── anomaly_detection.ipynb ├── instance_segmentation.ipynb ├── object_detection.ipynb └── semantic_segmentation.ipynb /README.md: -------------------------------------------------------------------------------- 1 | # TFDetection 2 | TFDetection(tfdet) is a detection toolbox based on Tensorflow2 and Keras. 3 | 4 | 5 | ## Overview of Components 6 |
7 | Detector 8 |
9 | 10 | 11 | 12 | 15 | 18 | 21 | 22 | 23 | 38 | 60 | 67 | 68 | 69 |
13 | Object Detection 14 | 16 | Segmentation 17 | 19 | Anomaly Detection 20 |
24 |
    25 |
  • Faster R-CNN(2015)
    • 26 |
  • RetinaNet(2017)
    • 27 |
  • YoloV3(2018)
    • 28 |
  • YoloV3 Tiny(2018)
    • 29 |
  • Cascade R-CNN(2018)
    • 30 |
  • FCOS(2019)
    • 31 |
  • Hybrid Task Cascade(2019)
    • 32 |
  • EfficientDet(2019)
    • 33 |
  • YoloV4(2020)
    • 34 |
  • YoloV4 Tiny(2020)
    • 35 |
  • EfficientDet Lite(2020)
    • 36 |
37 | 		39 |
    40 |
  • Semantic Segmentation
    • 41 |
      42 |
    • FCN(2015)
      • 43 |
    • UNet(2016)
      • 44 |
    • PSPNet(2017)
      • 45 |
    • DeepLab V3(2017)
      • 46 |
    • DeepLab V3+(2018)
      • 47 |
    • UNet++(2018)
      • 48 |
    • UperNet(2018)
      • 49 |
    50 |
51 |
    52 |
  • Instance Segmentation
    • 53 |
      54 |
    • Mask R-CNN(2017)
      • 55 |
    • Cascade R-CNN(2018)
      • 56 |
    • Hybrid Task Cascade(2019)
      • 57 |
    58 |
59 | 		61 |
    62 |
  • SPADE(2020)
    • 63 |
  • PaDiM(2020)
    • 64 |
  • PatchCore(2021)
    • 65 |
66 |
70 | 71 |
72 | Components 73 |
74 | 75 | 76 | 77 | 80 | 83 | 86 | 87 | 88 | 108 | 115 | 156 | 157 | 158 |
78 | Backbone 79 | 81 | Neck 82 | 84 | Other 85 |
89 |
    90 |
  • VGGNet(2015)
    • 91 |
  • ResNet(2016)
    • 92 |
  • Wide ResNet(2016)
    • 93 |
  • DenseNet(2016)
    • 94 |
  • ResNeXt(2016)
    • 95 |
  • MobileNetV2(2018)
    • 96 |
  • DarkNet(2018)
    • 97 |
  • MobileNetV3(2019)
    • 98 |
  • EfficientNet(2019)
    • 99 |
  • EfficientNet Lite(2020)
    • 100 |
  • ResNeSt(2020)
    • 101 |
  • CSP DarkNet(2020)
    • 102 |
  • SwinTransformer(2021)
    • 103 |
  • EfficientNetV2(2021)
    • 104 |
  • SwinTransformerV2(2021)
    • 105 |
  • ConvNeXt(2022)
    • 106 |
107 | 		109 |
    110 |
  • FPN(2017)
    • 111 |
  • PANet(2018)
    • 112 |
  • BiFPN(2019)
    • 113 |
114 | 		116 |
    117 |
  • Assign
    • 118 |
      119 |
    • Max Iou
      • 120 |
    • Center Region
      • 121 |
    • Point
      • 122 |
    • ATSS
      • 123 |
    • Sim OTA
      • 124 |
    • Align OTA
      • 125 |
    126 |
127 |
    128 |
  • Augmentation
    • 129 |
      130 |
    • Albumentations
      • 131 |
    • Mosaic
      • 132 |
    • CutMix
      • 133 |
    • CutOut
      • 134 |
    • MixUp
      • 135 |
    • CopyPaste
      • 136 |
    • MMDetction
      • 137 |
    • Yolo
      • 138 |
    139 |
140 |
    141 |
  • Metric
    • 142 |
      143 |
    • Mean Average Precision
      • 144 |
    • Mean IoU
      • 145 |
    146 |
147 |
    148 |
  • Export
    • 149 |
      150 |
    • TFLite
      • 151 |
    • ONNX
      • 152 |
    • TensorRT
      • 153 |
    154 |
155 |
159 | 160 | 161 | ## Requirements 162 | 163 | - python 3.8▲ 164 | - tensorflow 2.4▲ (If you installed tensorflow latest, more backbone is available.) 165 | - opencv-python 166 | 167 | 168 | ## Contributor 169 | 170 | * Hyungjin Kim(flslzk@gmail.com) 171 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | def readme(): 4 | with open("./README.md", encoding = "utf-8") as file: 5 | content = file.read() 6 | return content 7 | 8 | setuptools.setup(name = "tfdet", 9 | version = "1.0.0", 10 | author = "Hyungjin Kim", 11 | author_email = "flslzk@gmail.com", 12 | url = "https://github.com/Burf/tfdetection", 13 | description = "A Detection Toolbox for Tensorflow2", 14 | long_description = readme(), 15 | long_description_content_type = "text/markdown", 16 | license = "Apache License 2.0", 17 | install_requires = ["tensorflow>=2.4", "opencv-python"], 18 | packages = setuptools.find_packages(), 19 | zip_safe = False) 20 | -------------------------------------------------------------------------------- /tfdet/__init__.py: -------------------------------------------------------------------------------- 1 | from . import builder 2 | from . import callback 3 | from . import core 4 | from . import dataset 5 | from . import export 6 | from . import model 7 | from . import util -------------------------------------------------------------------------------- /tfdet/builder/__init__.py: -------------------------------------------------------------------------------- 1 | from .api import build_transform 2 | from .util import walk_module, find_module -------------------------------------------------------------------------------- /tfdet/builder/api.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import inspect 3 | 4 | import tfdet 5 | from .util import walk_module, find_module 6 | 7 | def build_transform(transform = None, key = "name"): 8 | """ 9 | transform = {'name':transform name or func, **kwargs} or transform name or func #find module in tfdet.dataset.transform and map kwargs. 10 | kwargs["sample_size"] > Covnert transform into multi_transform.(If transform doesn't need sample_size.) 11 | """ 12 | store = walk_module(tfdet.dataset.transform) 13 | 14 | result = [] 15 | for info in ([transform] if not isinstance(transform, (tuple, list)) else transform): 16 | if info is not None: 17 | if isinstance(info, dict): 18 | if "name" in info: 19 | info = info.copy() 20 | func = info.pop("name") 21 | if isinstance(func, str): 22 | func = find_module(store, func) 23 | if callable(func): 24 | func_spec = inspect.getfullargspec(func) 25 | keys = func_spec.args + func_spec.kwonlyargs 26 | if "sample_size" not in keys and "sample_size" in info: 27 | func = tfdet.dataset.multi_transform(func, info.pop("sample_size")) 28 | func = functools.partial(func, **{k:v for k, v in info.items() if k in keys}) 29 | result.append(func) 30 | elif isinstance(info, str): 31 | result.append(find_module(store, info)) 32 | elif callable(info): 33 | result.append(info) 34 | 35 | if len(result) == 0: 36 | reuslt = None 37 | elif len(result) == 1: 38 | result = result[0] 39 | return result -------------------------------------------------------------------------------- /tfdet/builder/util.py: -------------------------------------------------------------------------------- 1 | import types 2 | def walk_module(module, depth = 0, items = [types.FunctionType, type]): #type > class type 3 | store = {} 4 | for name in dir(module): 5 | attr = getattr(module, name) 6 | if isinstance(attr, types.ModuleType): 7 | if 0 < depth: 8 | r = walk_modules(attr, depth - 1, items = items) 9 | store.update(r) 10 | else: 11 | for item in items: 12 | if isinstance(attr, item): 13 | store["{0}.{1}".format(module.__name__, name)] = attr 14 | break 15 | return store 16 | 17 | def find_module(store, name): 18 | result = None 19 | for k, v in store.items(): 20 | if name == k[-len(name):]: 21 | result = v 22 | break 23 | return result -------------------------------------------------------------------------------- /tfdet/callback/__init__.py: -------------------------------------------------------------------------------- 1 | from .scheduler import (LearningRateScheduler, LearningRateSchedulerStep, 2 | WarmUpLearningRateScheduler, LinearLearningRateScheduler, CosineLearningRateScheduler, 3 | WarmUpLinearLearningRateScheduler, WarmUpCosineLearningRateScheduler, 4 | WarmUpLearningRateSchedulerStep, LinearLearningRateSchedulerStep, CosineLearningRateSchedulerStep, 5 | WarmUpLinearLearningRateSchedulerStep, WarmUpCosineLearningRateSchedulerStep) 6 | from .metric import MeanAveragePrecision, CoCoMeanAveragePrecision, MeanIoU 7 | from .util import EMA -------------------------------------------------------------------------------- /tfdet/callback/util.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from tfdet.core.util import EMA as ema_util 5 | 6 | class EMA(tf.keras.callbacks.Callback): 7 | """ 8 | 1) ema = EMA(decay = 0.9999, auto_apply = False) #auto_apply > ema.apply() in on_epoch_end, ema.restore() in on_epoch_begin 9 | 2) model.fit(..., 10 | callbacks=[..., 11 | ema, 12 | ...]) 13 | 14 | if step is None, update by epoch. else update by N step 15 | recommend step > max(round(64 / batch_size), 1) 16 | """ 17 | def __init__(self, decay = 0.9999, auto_apply = True, step = None, warm_up_epoch = 0, n_update = 0, ramp = 2000, init_model = None, apply_model = None, name = "ema", **kwargs): 18 | super(EMA, self).__init__(**kwargs) 19 | self.decay = decay 20 | self.auto_apply = auto_apply 21 | self.step = step 22 | self.warm_up_epoch = warm_up_epoch 23 | self.n_update = n_update 24 | self.ramp = ramp 25 | self.apply_model = apply_model 26 | self.name = name 27 | 28 | self.ema = None 29 | self.model = None 30 | self._step = step 31 | self.step_count = 0 32 | 33 | if init_model is not None: 34 | self.ema = ema_util(init_model, decay = self.decay, n_update = self.n_update, ramp = self.ramp) 35 | 36 | def on_train_begin(self, logs = None): 37 | if self.ema is None: 38 | self.ema = ema_util(self.model, decay = self.decay, n_update = self.n_update, ramp = self.ramp) 39 | self.ema.model = self.model 40 | 41 | def on_epoch_begin(self, epoch, logs = None): 42 | if self.step is not None: 43 | if self.warm_up_epoch <= epoch: 44 | self._step = self.step 45 | else: 46 | #https://github.com/WongKinYiu/yolov7 47 | self._step = max(1, round(np.interp(epoch, [0, self.warm_up_epoch - 1], [1, self.step]))) 48 | if self.auto_apply: 49 | self.restore(self.apply_model if isinstance(self.apply_model, tf.keras.Model) else self.model) 50 | 51 | def on_epoch_end(self, epoch, logs = {}): 52 | if self._step is None: 53 | self.ema.update(self.model) 54 | if self.auto_apply: 55 | self.apply(self.apply_model if isinstance(self.apply_model, tf.keras.Model) else self.model) 56 | logs["{0}_n_update".format(self.name)] = self.ema.n_update 57 | 58 | def on_train_batch_begin(self, step, logs = None): 59 | self.step_count += 1 60 | 61 | def on_train_batch_end(self, step, logs = None): 62 | if self._step is not None and round(self.step_count % self._step) == 0: 63 | self.ema.update(self.model) 64 | 65 | def apply(self, model = None): 66 | if self.ema is not None: 67 | self.ema.apply(model) 68 | 69 | def restore(self, model = None): 70 | if self.ema is not None: 71 | self.ema.restore(model) 72 | 73 | @staticmethod 74 | def get_n_update(init_epoch, step = None, total_step = None, warm_up_epoch = 0): 75 | if step is None: 76 | n_update = init_epoch 77 | else: 78 | n_update = 0 79 | step_count = 0 80 | for epoch in range(init_epoch): 81 | if warm_up_epoch <= epoch: 82 | _step = step 83 | else: 84 | _step = max(1, round(np.interp(epoch, [0, warm_up_epoch - 1], [1, step]))) 85 | for i in range(total_step): 86 | step_count += 1 87 | if round(step_count % _step) == 0: 88 | n_update += 1 89 | return n_update -------------------------------------------------------------------------------- /tfdet/core/__init__.py: -------------------------------------------------------------------------------- 1 | from . import anchor 2 | from . import assign 3 | from . import bbox 4 | from . import loss 5 | from . import metric 6 | from . import ops 7 | from . import util -------------------------------------------------------------------------------- /tfdet/core/anchor/__init__.py: -------------------------------------------------------------------------------- 1 | from .generator import generate_anchors, generate_yolo_anchors, generate_points 2 | from .util import generate_hist_scale, generate_uniform_scale, generate_kmeans_scale -------------------------------------------------------------------------------- /tfdet/core/anchor/util.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from ..bbox import overlap_bbox_numpy as overlap_bbox 5 | 6 | def generate_hist_scale(bbox_true, count = 5, decimal = 4): 7 | bbox_true = tf.reshape(bbox_true, (-1, 4)) 8 | indices = tf.where(tf.reduce_max(tf.cast(0 < bbox_true, tf.int32), axis = -1))[:, 0] 9 | bbox_true = tf.gather(bbox_true, indices) 10 | 11 | h = bbox_true[:, 3] - bbox_true[:, 1] 12 | w = bbox_true[:, 2] - bbox_true[:, 0] 13 | h, w = np.histogram2d(h, w, bins = count)[1:] 14 | scale = np.stack([h, w], axis = -1) 15 | scale = np.mean([scale[:count], scale[-count:]], axis = 0) 16 | return np.sort(np.round(scale, decimal), axis = 0) 17 | 18 | def generate_uniform_scale(min = 0.03125, max = 0.5, count = 5): 19 | return [min + (max - min) / (count - 1) * index for index in range(count)] 20 | 21 | def generate_kmeans_scale(bbox_true, k = 5, decimal = 4, method = np.median, mode = "normal"): 22 | bbox_true = np.reshape(bbox_true, [-1, 4]).astype(np.float32) 23 | bbox_true = bbox_true[np.max(0 < bbox_true, axis = -1)] 24 | wh = bbox_true - np.tile(bbox_true[..., :2], 2) #x1, y1, x2, y2 -> 0, 0, w, h 25 | 26 | n_bbox = len(wh) 27 | last_nearest = np.zeros((n_bbox,)) 28 | 29 | k = min(n_bbox, k) 30 | cluster = wh[np.random.choice(n_bbox, k, replace = False)] 31 | while True: 32 | overlaps = np.transpose(overlap_bbox(cluster, wh, mode = mode)) #(n_bbox, k) 33 | cur_nearest = np.argmin(1 - overlaps, axis = 1) 34 | if np.all(last_nearest == cur_nearest): 35 | break 36 | for index in range(k): 37 | target_wh = wh[cur_nearest == index] 38 | cluster[index] = method(target_wh, axis = 0) if 0 < len(target_wh) else cluster[index] 39 | last_nearest = cur_nearest 40 | return np.sort(np.round(cluster[..., 2:], decimal), axis = 0) 41 | -------------------------------------------------------------------------------- /tfdet/core/assign/__init__.py: -------------------------------------------------------------------------------- 1 | from .atss import atss 2 | from .center_region import center_region 3 | from .max_iou import max_iou 4 | from .ota import sim_ota, align_ota 5 | from .point import point 6 | 7 | from .sampler import random_sampler -------------------------------------------------------------------------------- /tfdet/core/assign/atss.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..bbox import overlap_bbox, isin 4 | from ..ops import euclidean_matrix 5 | 6 | def atss(y_true, bbox_true, y_pred, bbox_pred, k = 9, threshold = 0.01, min_threshold = 0.0001, extra_length = None, mode = "normal"): 7 | #https://arxiv.org/abs/1912.02424 8 | k = tf.minimum(k, tf.shape(bbox_pred)[0]) 9 | overlaps = overlap_bbox(bbox_true, bbox_pred, mode = mode) #(T, P) 10 | dist = euclidean_matrix(bbox_true, bbox_pred) #(T, P) 11 | sort_indices = tf.argsort(dist, axis = -1) #(T, P) 12 | candidate_indices = sort_indices[..., :k] #(T, K) 13 | 14 | candidate_overlaps = tf.gather(overlaps, candidate_indices, batch_dims = -1) #(T, K) 15 | candidate_threshold = tf.reduce_mean(candidate_overlaps, axis = -1) + tf.math.reduce_std(candidate_overlaps, axis = -1) 16 | candidate_flag = tf.greater_equal(candidate_overlaps, tf.expand_dims(candidate_threshold, axis = -1)) 17 | 18 | true_count = tf.shape(bbox_true)[0] 19 | pred_count = tf.shape(bbox_pred)[0] 20 | candidate_flag = tf.concat([candidate_flag, tf.zeros((true_count, pred_count - k), dtype = tf.bool)], axis = -1) #(T, K) + (T, P - K) 21 | candidate_flag = tf.gather(candidate_flag, tf.argsort(sort_indices, axis = -1), batch_dims = -1) #(T, P) 22 | 23 | isin_flag = isin(bbox_true, bbox_pred, extra_length = extra_length, mode = "rect") #(T, P) 24 | overlaps = tf.where(tf.logical_and(candidate_flag, isin_flag), overlaps, 0) #(T, P) 25 | overlaps = tf.transpose(overlaps) #(P, T) 26 | 27 | max_iou = tf.reduce_max(overlaps, axis = -1) 28 | match = tf.where(max(threshold, min_threshold) <= max_iou, 1, -1) 29 | 30 | positive_indices = tf.where(match == 1)[:, 0] 31 | negative_indices = tf.where(match == -1)[:, 0] 32 | 33 | positive_overlaps = tf.gather(overlaps, positive_indices) 34 | true_indices = tf.cond(tf.greater(tf.shape(positive_overlaps)[1], 0), true_fn = lambda: tf.argmax(positive_overlaps, axis = -1), false_fn = lambda: tf.cast(tf.constant([]), tf.int64)) 35 | return true_indices, positive_indices, negative_indices -------------------------------------------------------------------------------- /tfdet/core/assign/center_region.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..bbox import overlap_bbox, scale_bbox, isin 4 | 5 | def center_region(y_true, bbox_true, y_pred, bbox_pred, positive_scale = 0.2, negative_scale = 0.5, threshold = 0.01, min_threshold = 0.0001, extra_length = None, mode = "normal"): 6 | #https://arxiv.org/abs/1901.03278 7 | pos_bbox_true = scale_bbox(bbox_true, positive_scale) 8 | neg_bbox_true = scale_bbox(bbox_true, negative_scale) 9 | 10 | pos_flag = tf.transpose(isin(pos_bbox_true, bbox_pred, extra_length = extra_length, mode = "rect")) #(P, T) 11 | neg_flag = tf.transpose(~isin(neg_bbox_true, bbox_pred, extra_length = extra_length, mode = "rect")) #(P, T) 12 | #ignore_flag = tf.logical_and(~pos_flag, ~neg_flag) 13 | overlaps = overlap_bbox(bbox_pred, bbox_true, mode = mode) #(P, T) 14 | overlaps = tf.where(pos_flag, overlaps, 0) 15 | neg_overlaps = tf.where(neg_flag, -1, 0) 16 | 17 | max_iou = tf.reduce_max(overlaps, axis = -1) 18 | match = tf.reduce_min(neg_overlaps, axis = -1) 19 | match = tf.where(max(threshold, min_threshold) <= max_iou, 1, match) 20 | 21 | positive_indices = tf.where(match == 1)[:, 0] 22 | negative_indices = tf.where(match == -1)[:, 0] 23 | 24 | positive_overlaps = tf.gather(overlaps, positive_indices) 25 | true_indices = tf.cond(tf.greater(tf.shape(positive_overlaps)[1], 0), true_fn = lambda: tf.argmax(positive_overlaps, axis = -1), false_fn = lambda: tf.cast(tf.constant([]), tf.int64)) 26 | return true_indices, positive_indices, negative_indices -------------------------------------------------------------------------------- /tfdet/core/assign/max_iou.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..bbox import overlap_bbox 4 | 5 | def max_iou(y_true, bbox_true, y_pred, bbox_pred, positive_threshold = 0.5, negative_threshold = 0.4, min_threshold = 0.0001, match_low_quality = True, mode = "normal"): 6 | overlaps = overlap_bbox(bbox_pred, bbox_true, mode = mode) #(P, T) 7 | max_iou = tf.reduce_max(overlaps, axis = -1) 8 | 9 | match = tf.where(max_iou < negative_threshold, -1, 0) 10 | match = tf.where(max(positive_threshold, min_threshold) <= max_iou, 1, match) 11 | if match_low_quality: 12 | max_gt_iou = tf.reduce_max(overlaps, axis = 0, keepdims = True) 13 | max_gt_iou = tf.where(min_threshold <= max_gt_iou, max_gt_iou, -1) 14 | low_match = tf.reduce_max(tf.cast(overlaps == max_gt_iou, match.dtype), axis = -1) 15 | match = tf.where(low_match == 1, low_match, match) 16 | 17 | positive_indices = tf.where(match == 1)[:, 0] 18 | negative_indices = tf.where(match == -1)[:, 0] 19 | 20 | positive_overlaps = tf.gather(overlaps, positive_indices) 21 | true_indices = tf.cond(tf.greater(tf.shape(positive_overlaps)[1], 0), true_fn = lambda: tf.argmax(positive_overlaps, axis = -1), false_fn = lambda: tf.cast(tf.constant([]), tf.int64)) 22 | return true_indices, positive_indices, negative_indices -------------------------------------------------------------------------------- /tfdet/core/assign/point.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..bbox import overlap_point 4 | 5 | def point(y_true, bbox_true, y_pred, point_pred, regress_range = None, threshold = 0.0001, min_threshold = 0.0001): 6 | overlaps = tf.transpose(overlap_point(bbox_true, point_pred, regress_range)) #(P, T) 7 | max_area = tf.reduce_max(overlaps, axis = -1) 8 | match = tf.where(max(threshold, min_threshold) <= max_area, 1, -1) 9 | 10 | positive_indices = tf.where(match == 1)[:, 0] 11 | negative_indices = tf.where(match == -1)[:, 0] 12 | 13 | positive_overlaps = tf.gather(overlaps, positive_indices) 14 | true_indices = tf.cond(tf.greater(tf.shape(positive_overlaps)[1], 0), true_fn = lambda: tf.argmax(positive_overlaps, axis = -1), false_fn = lambda: tf.cast(tf.constant([]), tf.int64)) 15 | return true_indices, positive_indices, negative_indices -------------------------------------------------------------------------------- /tfdet/core/assign/sampler.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def random_sampler(true_indices, positive_indices, negative_indices, sampling_count = 256, positive_ratio = 0.5, return_count = False): 4 | if isinstance(sampling_count, int) and 0 < sampling_count: 5 | positive_count = tf.cast(sampling_count * positive_ratio, tf.int32) 6 | indices = tf.range(tf.shape(positive_indices)[0]) 7 | indices = tf.random.shuffle(indices)[:positive_count] 8 | positive_indices = tf.gather(positive_indices, indices) 9 | true_indices = tf.gather(true_indices, indices) 10 | positive_count = tf.cast(tf.shape(positive_indices)[0], tf.float32) 11 | negative_count = tf.cast(1 / positive_ratio * positive_count - positive_count, tf.int32) 12 | negative_indices = tf.random.shuffle(negative_indices)[:negative_count] 13 | if return_count: 14 | return true_indices, positive_indices, negative_indices, sampling_count 15 | else: 16 | return true_indices, positive_indices, negative_indices -------------------------------------------------------------------------------- /tfdet/core/bbox/__init__.py: -------------------------------------------------------------------------------- 1 | from .coder import bbox2delta, delta2bbox, bbox2yolo, yolo2bbox, bbox2offset, offset2bbox, offset2centerness 2 | from .overlap import overlap_bbox, overlap_point, overlap_bbox_numpy 3 | from .util import scale_bbox, isin, iou, iou_numpy, random_bbox -------------------------------------------------------------------------------- /tfdet/core/bbox/coder.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | def bbox2delta(bbox_true, bbox_pred, mean = [0., 0., 0., 0.], std = [1., 1., 1., 1.]): 5 | true_h = bbox_true[..., 3] - bbox_true[..., 1] 6 | true_w = bbox_true[..., 2] - bbox_true[..., 0] 7 | true_center_x = bbox_true[..., 0] + 0.5 * true_w 8 | true_center_y = bbox_true[..., 1] + 0.5 * true_h 9 | 10 | pred_h = bbox_pred[..., 3] - bbox_pred[..., 1] 11 | pred_w = bbox_pred[..., 2] - bbox_pred[..., 0] 12 | pred_h = tf.maximum(pred_h, tf.keras.backend.epsilon()) #tf.where(pred_h <= 0, tf.keras.backend.epsilon(), pred_h) 13 | pred_w = tf.maximum(pred_w, tf.keras.backend.epsilon()) #tf.where(pred_w <= 0, tf.keras.backend.epsilon(), pred_w) 14 | pred_center_x = bbox_pred[..., 0] + 0.5 * pred_w 15 | pred_center_y = bbox_pred[..., 1] + 0.5 * pred_h 16 | 17 | x = (true_center_x - pred_center_x) / pred_w 18 | y = (true_center_y - pred_center_y) / pred_h 19 | w = true_w / pred_w 20 | h = true_h / pred_h 21 | w = tf.math.log(tf.maximum(w, tf.keras.backend.epsilon())) #tf.where(w <= 0, tf.keras.backend.epsilon(), w) 22 | h = tf.math.log(tf.maximum(h, tf.keras.backend.epsilon())) #tf.where(h <= 0, tf.keras.backend.epsilon(), h) 23 | 24 | delta = tf.stack([x, y, w, h], axis = -1) 25 | if mean is not None: 26 | delta = delta - tf.cast(mean, delta.dtype) 27 | if std is not None: 28 | delta = delta / tf.maximum(tf.cast(std, delta.dtype), tf.keras.backend.epsilon()) 29 | return delta 30 | 31 | def delta2bbox(bbox, delta, mean = [0., 0., 0., 0.], std = [1., 1., 1., 1.], clip_ratio = 16 / 1000): 32 | if std is not None: 33 | delta = delta * tf.cast(std, delta.dtype) 34 | if mean is not None: 35 | delta = delta + tf.cast(mean, delta.dtype) 36 | h = bbox[..., 3] - bbox[..., 1] 37 | w = bbox[..., 2] - bbox[..., 0] 38 | center_x = bbox[..., 0] + 0.5 * w 39 | center_y = bbox[..., 1] + 0.5 * h 40 | center_x += delta[..., 0] * w 41 | center_y += delta[..., 1] * h 42 | delta_h = delta[..., 3] 43 | delta_w = delta[..., 2] 44 | if isinstance(clip_ratio, float): 45 | clip_value = np.abs(np.log(clip_ratio)) 46 | clip_value = tf.cast(clip_value, delta_h.dtype) 47 | delta_h = tf.clip_by_value(delta_h, -clip_value, clip_value) 48 | delta_w = tf.clip_by_value(delta_w, -clip_value, clip_value) 49 | h *= tf.exp(delta_h) 50 | w *= tf.exp(delta_w) 51 | x1 = center_x - 0.5 * w 52 | y1 = center_y - 0.5 * h 53 | x2 = x1 + w 54 | y2 = y1 + h 55 | bbox = tf.stack([x1, y1, x2, y2], axis = -1) 56 | return bbox 57 | 58 | def bbox2yolo(bbox_true, bbox_pred): 59 | pred_h = bbox_pred[..., 3] - bbox_pred[..., 1] 60 | pred_w = bbox_pred[..., 2] - bbox_pred[..., 0] 61 | pred_h = tf.maximum(pred_h, tf.keras.backend.epsilon()) #tf.where(pred_h <= 0, tf.keras.backend.epsilon(), pred_h) 62 | pred_w = tf.maximum(pred_w, tf.keras.backend.epsilon()) #tf.where(pred_w <= 0, tf.keras.backend.epsilon(), pred_w) 63 | 64 | x = (bbox_true[..., 0] - bbox_pred[..., 0]) / pred_w 65 | y = (bbox_true[..., 1] - bbox_pred[..., 1]) / pred_h 66 | w = (bbox_true[..., 2] - bbox_true[..., 0]) / pred_w 67 | h = (bbox_true[..., 3] - bbox_true[..., 1]) / pred_h 68 | w = tf.math.log(tf.maximum(w, tf.keras.backend.epsilon())) #tf.where(w <= 0, tf.keras.backend.epsilon(), w) 69 | h = tf.math.log(tf.maximum(h, tf.keras.backend.epsilon())) #tf.where(h <= 0, tf.keras.backend.epsilon(), h) 70 | 71 | delta = tf.stack([x, y, w, h], axis = -1) 72 | return delta 73 | 74 | def yolo2bbox(bbox, delta, clip_ratio = 16 / 1000): 75 | h = bbox[..., 3] - bbox[..., 1] 76 | w = bbox[..., 2] - bbox[..., 0] 77 | x1 = bbox[..., 0] + delta[..., 0] * w 78 | y1 = bbox[..., 1] + delta[..., 1] * h 79 | delta_h = delta[..., 3] 80 | delta_w = delta[..., 2] 81 | if isinstance(clip_ratio, float): 82 | clip_value = np.abs(np.log(clip_ratio)) 83 | clip_value = tf.cast(clip_value, delta_h.dtype) 84 | delta_h = tf.clip_by_value(delta_h, -clip_value, clip_value) 85 | delta_w = tf.clip_by_value(delta_w, -clip_value, clip_value) 86 | h *= tf.exp(delta_h) 87 | w *= tf.exp(delta_w) 88 | x2 = x1 + w 89 | y2 = y1 + h 90 | bbox = tf.stack([x1, y1, x2, y2], axis = -1) 91 | return bbox 92 | 93 | def bbox2offset(bbox_true, points): 94 | x1, y1, x2, y2 = tf.split(bbox_true, 4, axis = -1) 95 | px, py = tf.split(points, 2, axis = -1) 96 | offset = tf.concat([px - x1, py - y1, x2 - px, y2 - py], axis = -1) #left, top, right, bottom 97 | return offset 98 | 99 | def offset2bbox(points, offset): 100 | x, y = tf.split(points, 2, axis = -1) 101 | left, top, right, bottom = tf.split(offset, 4, axis = -1) 102 | bbox = tf.concat([x - left, y - top, x + right, y + bottom], axis = -1) 103 | return bbox 104 | 105 | def offset2centerness(offset): 106 | left, top, right, bottom = tf.split(offset, 4, axis = -1) 107 | lr = tf.concat([left, right], axis = -1) 108 | tb = tf.concat([top, bottom], axis = -1) 109 | max_lr = tf.reduce_max(lr, axis = -1, keepdims = True) 110 | max_tb = tf.reduce_max(tb, axis = -1, keepdims = True) 111 | centerness = tf.sqrt((tf.reduce_min(lr, axis = -1, keepdims = True) / tf.maximum(max_lr, tf.keras.backend.epsilon())) * (tf.reduce_min(tb, axis = -1, keepdims = True) / tf.maximum(max_tb, tf.keras.backend.epsilon()))) 112 | return centerness -------------------------------------------------------------------------------- /tfdet/core/bbox/overlap.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from .util import iou, iou_numpy 5 | 6 | def overlap_bbox(bbox_true, bbox_pred, mode = "normal"): 7 | """ 8 | bbox_true = [[x1, y1, x2, y2], ...] #(N, bbox) 9 | bbox_pred = [[x1, y1, x2, y2], ...] #(M, bbox) 10 | 11 | overlaps = true & pred iou matrix #(N, M) 12 | """ 13 | if mode not in ("normal", "foreground", "general", "complete", "distance"): 14 | raise ValueError("unknown mode '{0}'".format(mode)) 15 | 16 | true_count = tf.shape(bbox_true)[0] 17 | pred_count = tf.shape(bbox_pred)[0] 18 | 19 | bbox_true = tf.reshape(tf.tile(tf.expand_dims(bbox_true, 0), [1, 1, pred_count]), [-1, 4]) 20 | bbox_pred = tf.tile(bbox_pred, [true_count, 1]) 21 | 22 | overlaps = iou(bbox_true, bbox_pred, mode = mode) 23 | overlaps = tf.reshape(overlaps, [true_count, pred_count]) 24 | return overlaps 25 | 26 | def overlap_point(bbox_true, points, regress_range = None): 27 | true_count = tf.shape(bbox_true)[0] 28 | pred_count = tf.shape(points)[0] 29 | 30 | bbox_true = tf.reshape(tf.tile(tf.expand_dims(bbox_true, 0), [1, 1, pred_count]), [-1, 4]) 31 | points = tf.tile(points, [true_count, 1]) 32 | 33 | x1, y1, x2, y2 = tf.split(bbox_true, 4, axis = -1) 34 | px, py = tf.split(points, 2, axis = -1) 35 | area = tf.reshape((x2 - x1) * (y2 - y1), [true_count, pred_count]) 36 | offset = tf.concat([px - x1, py - y1, x2 - px, y2 - py], axis = -1) #left, top, right, bottom 37 | offset = tf.reshape(offset, [true_count, pred_count, 4]) 38 | min_offset = tf.reduce_min(offset, axis = -1) 39 | 40 | overlap_flag = tf.greater(min_offset, 0) 41 | if regress_range is not None: 42 | max_offset = tf.reduce_max(offset, axis = -1) 43 | regress_range = tf.tile(regress_range, [true_count, 1]) 44 | regress_range = tf.reshape(regress_range, [true_count, pred_count, 2]) 45 | range_flag = tf.logical_and(tf.greater(max_offset, regress_range[..., 0]), tf.less_equal(max_offset, regress_range[..., 1])) 46 | overlap_flag = tf.logical_and(overlap_flag, range_flag) 47 | pad_area = tf.where(overlap_flag, area, tf.reduce_max(area) + 1) 48 | min_flag = tf.equal(area, tf.reduce_min(pad_area, axis = 0, keepdims = True)) 49 | overlaps = tf.where(min_flag, area, 0) 50 | return overlaps 51 | 52 | def overlap_bbox_numpy(bbox_true, bbox_pred, mode = "normal", e = 1e-12): 53 | """ 54 | bbox_true = [[x1, y1, x2, y2], ...] #(N, bbox) 55 | bbox_pred = [[x1, y1, x2, y2], ...] #(M, bbox) 56 | 57 | overlaps = true & pred iou matrix #(N, M) 58 | """ 59 | if mode not in ("normal", "foreground", "general", "complete", "distance"): 60 | raise ValueError("unknown mode '{0}'".format(mode)) 61 | 62 | true_count = np.shape(bbox_true)[0] 63 | pred_count = np.shape(bbox_pred)[0] 64 | 65 | bbox_true = np.reshape(np.tile(np.expand_dims(bbox_true, 0), [1, 1, pred_count]), [-1, 4]) 66 | bbox_pred = np.tile(bbox_pred, [true_count, 1]) 67 | 68 | overlaps = iou_numpy(bbox_true, bbox_pred, mode = mode, e = e) 69 | overlaps = np.reshape(overlaps, [true_count, pred_count]) 70 | return overlaps -------------------------------------------------------------------------------- /tfdet/core/loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .cross_entropy import binary_cross_entropy, categorical_cross_entropy, focal_binary_cross_entropy, focal_categorical_cross_entropy 2 | from .object_detection import iou, giou, diou, ciou 3 | from .regression import smooth_l1 4 | from .segmentation import dice, bce_dice, tversky, focal_tversky, iou_pixcel, generalized_dice, bce_generalized_dice 5 | from .util import regularize, weight_reduce_loss, resize_loss -------------------------------------------------------------------------------- /tfdet/core/loss/cross_entropy.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def binary_cross_entropy(y_true, y_pred, focal = False, alpha = .25, gamma = 1.5, weight = None, reduce = tf.reduce_mean): 4 | n_true_class = tf.shape(y_true)[-1] 5 | n_pred_class = tf.shape(y_pred)[-1] 6 | 7 | y_true = tf.cast(tf.cond(tf.logical_and(tf.equal(n_true_class, 1), tf.not_equal(n_pred_class, 1)), true_fn = lambda: tf.cast(tf.one_hot(tf.cast(y_true, tf.int32), n_pred_class)[..., 0, :], y_true.dtype), false_fn = lambda: y_true), y_pred.dtype) 8 | #y_pred = y_pred / (tf.reduce_sum(y_pred, axis = -1, keepdims = True) + tf.keras.backend.epsilon()) 9 | y_pred = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1 - tf.keras.backend.epsilon()) 10 | 11 | loss = -tf.stack([(1 - y_true) * tf.math.log(1 - y_pred), y_true * tf.math.log(y_pred)], axis = -1) 12 | if focal: 13 | alpha_factor = tf.ones_like(y_true) * alpha 14 | alpha_factor = tf.where(0.5 < y_true, alpha_factor, 1 - alpha_factor) 15 | focal_weight = tf.where(0.5 < y_true, 1 - y_pred, y_pred) 16 | focal_weight = alpha_factor * focal_weight ** gamma 17 | loss = tf.expand_dims(focal_weight, axis = -1) * loss 18 | loss = tf.reduce_sum(loss, axis = -1) 19 | if weight is not None: 20 | loss *= weight 21 | if reduce: 22 | axis = tf.range(tf.rank(y_true))[1:] 23 | loss = reduce(loss, axis = axis) 24 | loss = tf.reduce_mean(loss) 25 | return loss 26 | 27 | def categorical_cross_entropy(y_true, y_pred, focal = False, alpha = 1., gamma = 2., weight = None, reduce = tf.reduce_mean): 28 | n_true_class = tf.shape(y_true)[-1] 29 | n_pred_class = tf.shape(y_pred)[-1] 30 | 31 | y_true = tf.cast(tf.cond(tf.equal(n_true_class, 1), true_fn = lambda: tf.cast(tf.one_hot(tf.cast(y_true, tf.int32), n_pred_class)[..., 0, :], y_true.dtype), false_fn = lambda: y_true), y_pred.dtype) 32 | #y_pred = y_pred / (tf.reduce_sum(y_pred, axis = -1, keepdims = True) + tf.keras.backend.epsilon()) 33 | y_pred = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1 - tf.keras.backend.epsilon()) 34 | 35 | loss = -y_true * tf.math.log(y_pred) 36 | if focal: 37 | loss = alpha * tf.math.pow(1 - y_pred, gamma) * loss 38 | if weight is not None: 39 | loss *= weight 40 | loss = tf.reduce_sum(loss, axis = -1, keepdims = True) 41 | if reduce: 42 | axis = tf.range(tf.rank(y_true))[1:] 43 | loss = reduce(loss, axis = axis) 44 | loss = tf.reduce_mean(loss) 45 | return loss 46 | 47 | def focal_binary_cross_entropy(y_true, y_pred, alpha = .25, gamma = 1.5, weight = None, reduce = tf.reduce_mean): 48 | return binary_cross_entropy(y_true, y_pred, focal = True, alpha = alpha, gamma = gamma, weight = weight, reduce = reduce) 49 | 50 | def focal_categorical_cross_entropy(y_true, y_pred, alpha = 1., gamma = 2., weight = None, reduce = tf.reduce_mean): 51 | return categorical_cross_entropy(y_true, y_pred, focal = True, alpha = alpha, gamma = gamma, weight = weight, reduce = reduce) 52 | -------------------------------------------------------------------------------- /tfdet/core/loss/object_detection.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..bbox import iou as iou_calculator 4 | 5 | def iou(bbox_true, bbox_pred, reduce = tf.reduce_mean, mode = "normal"): 6 | iou = iou_calculator(bbox_pred, bbox_true, mode = mode) 7 | loss = 1 - iou 8 | if reduce: 9 | loss = reduce(loss) 10 | return loss 11 | 12 | def giou(bbox_true, bbox_pred, reduce = tf.reduce_mean, mode = "general"): 13 | return iou(bbox_true, bbox_pred, reduce = reduce, mode = mode) 14 | 15 | def ciou(bbox_true, bbox_pred, reduce = tf.reduce_mean, mode = "complete"): 16 | return iou(bbox_true, bbox_pred, reduce = reduce, mode = mode) 17 | 18 | def diou(bbox_true, bbox_pred, reduce = tf.reduce_mean, mode = "distance"): 19 | return iou(bbox_true, bbox_pred, reduce = reduce, mode = mode) -------------------------------------------------------------------------------- /tfdet/core/loss/regression.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def smooth_l1(y_true, y_pred, sigma = 3, reduce = tf.reduce_mean): 4 | diff = tf.abs(y_true - y_pred) 5 | less_than_one = tf.cast(tf.less(diff, (1.0 / sigma ** 2)), diff.dtype) 6 | loss = (less_than_one * (0.5 * sigma ** 2) * diff ** 2) + (1 - less_than_one) * (diff - (0.5 / sigma ** 2)) 7 | if reduce: 8 | loss = reduce(loss, axis = -1) 9 | loss = tf.reduce_mean(loss) 10 | return loss -------------------------------------------------------------------------------- /tfdet/core/loss/util.py: -------------------------------------------------------------------------------- 1 | import functools 2 | 3 | import tensorflow as tf 4 | 5 | def regularize(model, weight_decay = 1e-4, loss = tf.keras.regularizers.l2): 6 | weight_decay = weight_decay() if callable(weight_decay) else weight_decay 7 | reg_loss = [] 8 | for w in model.trainable_weights: 9 | if "gamma" not in w.name and "beta" not in w.name: 10 | l = loss(weight_decay)(w) 11 | reg_loss.append(l / tf.cast(tf.size(w), l.dtype)) 12 | return reg_loss 13 | 14 | def weight_reduce_loss(loss, weight = None, avg_factor = None): 15 | """ 16 | https://github.com/open-mmlab/mmdetection/blob/master/mmdet/models/losses/utils.py 17 | """ 18 | if weight is not None: 19 | loss = loss * weight 20 | if avg_factor is not None: 21 | loss = tf.reduce_sum(loss) / (avg_factor + tf.keras.backend.epsilon()) 22 | else: 23 | loss = tf.reduce_mean(loss) 24 | return loss 25 | 26 | def resize_loss(mask_true = None, mask_pred = None, loss = None, method = "bilinear", **kwargs): 27 | if callable(mask_true): 28 | return functools.partial(resize_loss, loss = mask_true, method = method, **kwargs) 29 | 30 | mask_pred = tf.image.resize(mask_pred, tf.shape(mask_true)[-3:-1], method = method) 31 | return loss(mask_true, mask_pred, **kwargs) -------------------------------------------------------------------------------- /tfdet/core/metric/__init__.py: -------------------------------------------------------------------------------- 1 | from .object_detection import MeanAveragePrecision, CoCoMeanAveragePrecision 2 | from .segmentation import MeanIoU -------------------------------------------------------------------------------- /tfdet/core/metric/segmentation.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from tfdet.core.util import metric2text 4 | 5 | class MeanIoU: 6 | def __init__(self, beta = 1, e = 1e-12, label = None, dtype = np.float32): 7 | """ 8 | run = MeanIoU()(*args) 9 | batch run = self.add(*batch_args) -> self.evaluate() 10 | """ 11 | self.beta = beta 12 | self.e = e 13 | self.label = label 14 | self.dtype = dtype 15 | 16 | self.reset() 17 | 18 | def reset(self): 19 | self.area_inter = None 20 | self.area_union = None 21 | self.area_true = None 22 | self.area_pred = None 23 | 24 | @property 25 | def iou(self): 26 | return self.evaluate(reduce = False) 27 | 28 | @property 29 | def mean_iou(self): 30 | return self.evaluate() 31 | 32 | @property 33 | def accuracy(self): 34 | return (self.area_inter / np.maximum(self.area_true, self.e)) if self.area_true is not None else [0.] 35 | 36 | @property 37 | def mean_accuracy(self): 38 | return np.mean(self.accuracy).item() 39 | 40 | @property 41 | def dice(self): 42 | return (2 * self.area_inter / np.maximum(self.area_true + self.area_pred, self.e)) if self.area_true is not None else [0.] 43 | 44 | @property 45 | def mean_dice(self): 46 | return np.mean(self.dice).item() 47 | 48 | @property 49 | def precision(self): 50 | return self.area_inter / np.maximum(self.area_pred, self.e) if self.area_true is not None else [0.] 51 | 52 | @property 53 | def recall(self): 54 | return self.accuracy 55 | 56 | @property 57 | def f1(self): 58 | return ((1 + self.beta**2) * np.multiply(self.precision, self.recall)) / np.maximum(np.add(np.multiply(self.beta**2, self.precision), self.recall), self.e) 59 | 60 | @property 61 | def mean_f1(self): 62 | return np.mean(self.f1).item() 63 | 64 | @property 65 | def summary(self): 66 | return [self.accuracy, self.iou, self.dice, self.f1] 67 | 68 | @property 69 | def summary_text(self): 70 | text = "" 71 | if self.area_true is not None: 72 | try: 73 | accuracy, iou, dice, f1 = self.summary 74 | info = {"accuracy":accuracy, "iou":iou, "dice":dice, "f1":f1} 75 | 76 | summary = [np.mean(accuracy).item(), np.mean(iou).item(), np.mean(dice).item(), np.mean(f1).item()] 77 | 78 | text = metric2text(info, summary = summary, label = self.label) 79 | except: 80 | pass 81 | return text 82 | 83 | def __call__(self, mask_true, mask_pred, reset = True): 84 | """ 85 | mask_true = #(batch_size, h, w, 1 or n_class) 86 | mask_pred = #(batch_size, h, w, n_class) 87 | """ 88 | if reset: 89 | self.reset() 90 | self.add(mask_true, mask_pred) 91 | return self.evaluate() 92 | 93 | def add(self, mask_true, mask_pred): 94 | """ 95 | mask_true = #(h, w, 1 or n_class) or (batch_size, h, w, 1 or n_class) 96 | mask_pred = #(h, w, n_class) or (batch_size, h, w, n_class) 97 | """ 98 | if np.ndim(mask_true) == 4: 99 | for index in range(len(mask_true)): 100 | self.add(mask_true[index], mask_pred[index]) 101 | else: 102 | n_class = np.shape(mask_pred)[-1] 103 | if self.area_true is None: 104 | self.area_inter = np.zeros(n_class, dtype = self.dtype) 105 | self.area_union = np.zeros(n_class, dtype = self.dtype) 106 | self.area_true = np.zeros(n_class, dtype = self.dtype) 107 | self.area_pred = np.zeros(n_class, dtype = self.dtype) 108 | 109 | if 1 < np.shape(mask_true)[-1]: 110 | mask_true = np.expand_dims(np.argmax(mask_true, axis = -1), axis = -1) 111 | mask_pred = np.expand_dims(np.argmax(mask_pred, axis = -1), axis = -1) 112 | 113 | inter = mask_pred[mask_pred == mask_true] 114 | area_inter = np.histogram(inter, bins = n_class, range = (0, n_class - 1))[0] 115 | area_true = np.histogram(mask_true, bins = n_class, range = (0, n_class - 1))[0] 116 | area_pred = np.histogram(mask_pred, bins = n_class, range = (0, n_class - 1))[0] 117 | if self.dtype is not None: 118 | area_inter = area_inter.astype(self.dtype) 119 | area_true = area_true.astype(self.dtype) 120 | area_pred = area_pred.astype(self.dtype) 121 | area_union = area_true + area_pred - area_inter 122 | self.area_inter += area_inter 123 | self.area_union += area_union 124 | self.area_true += area_true 125 | self.area_pred += area_pred 126 | 127 | def evaluate(self, reduce = True): 128 | if self.area_true is not None: 129 | iou = self.area_inter / np.maximum(self.area_union, self.e) 130 | if reduce: 131 | iou = np.mean(iou).item() 132 | return iou 133 | else: 134 | return 0. if reduce else [0.] 135 | -------------------------------------------------------------------------------- /tfdet/core/ops/__init__.py: -------------------------------------------------------------------------------- 1 | from .distance import * 2 | from .feature_extract import * 3 | from .initializer import * 4 | from .nms import * 5 | from .pooling import * 6 | from .roi_extract import * -------------------------------------------------------------------------------- /tfdet/core/ops/distance.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def mahalanobis(u, v, VI): 4 | delta = u - v 5 | out = tf.sqrt(tf.matmul(tf.matmul(tf.expand_dims(delta, axis = -2), VI), tf.expand_dims(delta, axis = -1))) 6 | return tf.squeeze(out, axis = -1) 7 | 8 | def euclidean(u, v): 9 | return tf.sqrt(tf.reduce_sum(tf.square(u - v), axis = -1)) 10 | 11 | def euclidean_matrix(u, v): 12 | u_norm = tf.reduce_sum(tf.square(u), axis = -1, keepdims = True) 13 | v_norm = tf.reduce_sum(tf.square(v), axis = -1, keepdims = True) 14 | dist = tf.sqrt(tf.maximum(tf.add(tf.transpose(v_norm) - 2 * tf.matmul(u, v, transpose_b = True), u_norm), 0)) 15 | return dist 16 | -------------------------------------------------------------------------------- /tfdet/core/ops/feature_extract.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | def feature_concat(x, y): 5 | b = tf.shape(x)[0] 6 | h1, w1, c1 = tf.keras.backend.int_shape(x)[1:] 7 | h2, w2, c2 = tf.keras.backend.int_shape(y)[1:] 8 | 9 | kernel_size = int(h1 / h2) 10 | x = tf.nn.space_to_depth(x, kernel_size) 11 | x = tf.reshape(x, [b, h2, w2, -1, c1]) 12 | shape = [1 if index != 2 else tf.shape(x)[-2] for index in range(tf.keras.backend.ndim(x), 0, -1)] 13 | y = tf.tile(tf.expand_dims(y, axis = -2), shape) 14 | out = tf.concat([x, y], axis = -1) 15 | out = tf.reshape(out, [b, h2, w2, -1]) 16 | out = tf.nn.depth_to_space(out, kernel_size) 17 | out = tf.reshape(out, [b, h1, w1, c1 + c2]) 18 | return out 19 | 20 | def feature_extract(feature, sampling_index = None, pool_size = 1, sub_sampling = False, concat = True, memory_reduce = False): 21 | if not isinstance(feature, list): 22 | feature = [feature] 23 | if isinstance(sampling_index, int): 24 | sampling_index = [sampling_index] 25 | feature = list(feature) 26 | 27 | if sub_sampling: 28 | feature.append(tf.nn.avg_pool(feature[-1], tf.keras.backend.int_shape(feature[-1])[-3], 1, padding = "VALID")) 29 | 30 | pad = 0 31 | for level in range(len(feature)): 32 | if 1 < pool_size: 33 | feature[level] = tf.nn.avg_pool(feature[level], pool_size, 1, padding = "SAME") 34 | if sampling_index is not None and (memory_reduce or not concat): 35 | h, w, c = tf.keras.backend.int_shape(feature[level])[1:] 36 | indices = tf.gather(sampling_index, tf.where(tf.logical_and(tf.greater_equal(sampling_index, pad), tf.less(sampling_index, c + pad))))[:, 0] - pad 37 | feature[level] = tf.reshape(tf.gather(feature[level], indices, axis = -1), [-1, h, w, len(indices)]) 38 | pad += c 39 | 40 | if concat: 41 | fv = feature[0] 42 | for f in feature[1:]: 43 | fv = feature_concat(fv, f) 44 | if sampling_index is not None and not memory_reduce: 45 | h, w, c = tf.keras.backend.int_shape(fv)[1:] 46 | fv = tf.reshape(tf.gather(fv, sampling_index, axis = -1), [-1, h, w, len(sampling_index)]) 47 | feature = [fv] 48 | 49 | if len(feature) == 1: 50 | feature = feature[0] 51 | return feature 52 | 53 | def core_sampling(*args, n_sample = 3, n_feature = "auto", eps = 0.9, index = False): 54 | try: 55 | from sklearn.random_projection import SparseRandomProjection, johnson_lindenstrauss_min_dim 56 | except Exception as e: 57 | print("If you want to use 'core_sampling', please install 'scikit-learn 0.14▲'") 58 | raise e 59 | if isinstance(n_sample, float): 60 | n_sample = int(len(args[0]) * n_sample) 61 | n_sample = max(min(n_sample, len(args[0])), 1 if len(args[0]) != 0 else 0) 62 | if n_feature == "auto": 63 | b, c = np.shape(args[0]) 64 | n_feature = max(min(johnson_lindenstrauss_min_dim(b, eps = eps), c), 1) 65 | m = SparseRandomProjection(n_components = n_feature, eps = eps) 66 | trans_data = m.fit_transform(args[0]) 67 | 68 | indices = [] 69 | min_dist = None 70 | target = np.expand_dims(trans_data[0], axis = 0) 71 | iter_range = range(n_sample) 72 | try: 73 | from tqdm import tqdm 74 | iter_range = tqdm(iter_range, total = n_sample, desc = "greedy sampling top-k center") 75 | except: 76 | pass 77 | for j in iter_range: #k center greedy 78 | dist = np.linalg.norm(trans_data - target, axis = -1, keepdims = True) 79 | min_dist = np.minimum(dist, min_dist) if min_dist is not None else dist 80 | min_index = np.argmax(min_dist) 81 | target = np.expand_dims(trans_data[min_index], axis = 0) 82 | min_dist[min_index] = 0 83 | indices.append(min_index) 84 | 85 | if not index: 86 | args = [(np.array(arg) if not isinstance(arg, np.ndarray) else arg)[indices] for arg in args] 87 | if len(args) == 1: 88 | args = args[0] 89 | return args 90 | else: 91 | return indices -------------------------------------------------------------------------------- /tfdet/core/ops/initializer.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | class PriorProbability(tf.keras.initializers.Initializer): 5 | """https://github.com/xuannianz/EfficientDet/blob/master/initializers.py""" 6 | def __init__(self, probability = 0.01): 7 | self.probability = probability 8 | 9 | def get_config(self): 10 | return {"probability":self.probability} 11 | 12 | def __call__(self, shape, dtype = None): 13 | # set bias to -log((1 - p)/p) for foreground 14 | result = np.ones(shape, dtype = np.float32) * -np.log((1 - self.probability) / self.probability) 15 | return result -------------------------------------------------------------------------------- /tfdet/core/ops/nms.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..bbox.coder import delta2bbox 4 | 5 | @tf.function 6 | def multiclass_nms(y_pred, bbox_pred, anchors = None, mask_pred = None, proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, 7 | ignore_label = 0, performance_count = 5000, coder_func = delta2bbox, **kwargs): 8 | """ 9 | y_pred = logit #(n_anchor, n_class) 10 | bbox_pred = delta #(n_anchor, 4) or n_class delta #(n_anchor, n_clss, 4) 11 | anchors = anchors #(n_anchor, 4) or points #(n_anchor, 2) 12 | mask_pred = mask #(n_anchor, H, W, 1 or n_class) 13 | 14 | y_pred = logit #(proposal_count, n_class) 15 | bbox_pred = normalized proposal [[x1, y1, x2, y2], ...] #(proposal_count, 4) 16 | mask_pred = mask #(proposal_count, H, W, 1) 17 | """ 18 | n_class = tf.keras.backend.int_shape(y_pred)[-1] 19 | score_threshold = [score_threshold] * n_class if isinstance(score_threshold, float) else score_threshold 20 | soft_nms_sigma = soft_nms 21 | if not isinstance(soft_nms, float): 22 | soft_nms_sigma = 0.5 if soft_nms else 0. 23 | if n_class == 1: 24 | ignore_label = None 25 | else: 26 | ignore_label = [ignore_label] if isinstance(ignore_label, int) else ignore_label 27 | keep_label = [i for i in range(n_class) if ignore_label is None or i not in ignore_label] 28 | 29 | #filtered by label 30 | if ignore_label is not None: 31 | label = tf.argmax(y_pred, axis = -1) 32 | flag = tf.zeros_like(label, dtype = tf.bool) 33 | for cls in ignore_label: 34 | flag = tf.logical_or(flag, label == cls) 35 | indices = tf.where(tf.logical_not(flag))[:, 0] 36 | 37 | y_pred = tf.gather(y_pred, indices) 38 | bbox_pred = tf.gather(bbox_pred, indices) 39 | if anchors is not None: 40 | anchors = tf.gather(anchors, indices) 41 | if mask_pred is not None: 42 | mask_pred = tf.gather(mask_pred, indices) 43 | 44 | #reduce by performance_count 45 | if isinstance(performance_count, int) and 0 < performance_count: 46 | top_indices = tf.nn.top_k(tf.reduce_max(y_pred, axis = -1), tf.minimum(performance_count, tf.shape(y_pred)[0]), sorted = True).indices 47 | y_pred = tf.gather(y_pred, top_indices) 48 | bbox_pred = tf.gather(bbox_pred, top_indices) 49 | if anchors is not None: 50 | anchors = tf.gather(anchors, top_indices) 51 | if mask_pred is not None: 52 | mask_pred = tf.gather(mask_pred, top_indices) 53 | 54 | bbox_flag = (tf.keras.backend.ndim(bbox_pred) == 3) 55 | mask_flag = (mask_pred is not None and tf.keras.backend.int_shape(mask_pred)[-1] != 1) 56 | if bbox_flag or mask_flag: 57 | if n_class == 1: 58 | label_indices = tf.zeros(tf.shape(y_pred)[0], dtype = tf.int32) 59 | else: 60 | label_indices = tf.argmax(y_pred, axis = -1, output_type = tf.int32) 61 | if bbox_flag: 62 | bbox_pred = tf.gather(bbox_pred, label_indices, batch_dims = 1) 63 | if mask_flag: 64 | mask_pred = tf.transpose(mask_pred, [0, 3, 1, 2]) 65 | mask_pred = tf.gather(mask_pred, label_indices, batch_dims = 1) 66 | mask_pred = tf.expand_dims(mask_pred, axis = -1) 67 | 68 | if anchors is not None and callable(coder_func): 69 | bbox_pred = coder_func(anchors, bbox_pred, **kwargs) 70 | bbox_pred = tf.clip_by_value(bbox_pred, 0, 1) 71 | x1, y1, x2, y2 = tf.split(bbox_pred, 4, axis = -1) 72 | bbox = tf.concat([y1, x1, y2, x2], axis = -1) 73 | 74 | dtype = bbox.dtype 75 | iou_threshold = tf.cast(iou_threshold, dtype) 76 | score_threshold = tf.cast(score_threshold, dtype) 77 | soft_nms_sigma = tf.cast(soft_nms_sigma, dtype) 78 | scores = [] 79 | indices = [] 80 | for cls in keep_label: 81 | threshold = score_threshold[cls] 82 | score = y_pred[..., cls] 83 | cls_indices = tf.image.non_max_suppression_with_scores(bbox, score, max_output_size = proposal_count, iou_threshold = iou_threshold, score_threshold = threshold, soft_nms_sigma = soft_nms_sigma)[0] 84 | scores.append(tf.gather(score, cls_indices)) 85 | indices.append(cls_indices) 86 | scores = tf.concat(scores, axis = 0) 87 | indices = tf.concat(indices, axis = 0) 88 | 89 | top_indices = tf.nn.top_k(scores, tf.minimum(proposal_count, tf.shape(scores)[0]), sorted = True).indices 90 | indices = tf.gather(indices, top_indices) 91 | y_pred = tf.gather(y_pred, indices) 92 | bbox_pred = tf.gather(bbox_pred, indices) 93 | if mask_pred is not None: 94 | h, w = tf.keras.backend.int_shape(mask_pred)[-3:-1] 95 | mask_pred = tf.gather(mask_pred, indices) 96 | 97 | pad_count = tf.maximum(proposal_count - tf.shape(bbox_pred)[0], 0) 98 | y_pred = tf.pad(y_pred, [[0, pad_count], [0, 0]]) 99 | bbox_pred = tf.pad(bbox_pred, [[0, pad_count], [0, 0]]) 100 | 101 | y_pred = tf.reshape(y_pred, [proposal_count, n_class]) 102 | bbox_pred = tf.reshape(bbox_pred, [proposal_count, 4]) 103 | result = y_pred, bbox_pred 104 | if mask_pred is not None: 105 | mask_pred = tf.pad(mask_pred, [[0, pad_count], [0, 0], [0, 0], [0, 0]]) 106 | mask_pred = tf.reshape(mask_pred, [proposal_count, h, w, 1]) 107 | result = y_pred, bbox_pred, mask_pred 108 | return result -------------------------------------------------------------------------------- /tfdet/core/ops/pooling.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | class AdaptivePooling(tf.keras.layers.Layer): 4 | #https://github.com/tensorflow/addons/blob/master/tensorflow_addons/layers/adaptive_pooling.py 5 | def __init__(self, scale, reduce = tf.reduce_mean, **kwargs): 6 | super(AdaptivePooling, self).__init__(**kwargs) 7 | self.scale = scale 8 | self.reduce = reduce 9 | 10 | def build(self, input_shape): 11 | self.loop = len(input_shape) - 2 12 | if not isinstance(self.scale, (tuple, list)): 13 | self.scale = [self.scale] * self.loop 14 | 15 | def call(self, inputs): 16 | out = inputs 17 | if 0 < self.loop: 18 | for i in range(self.loop): 19 | axis = 1 + 2 * i 20 | out = tf.split(out, self.scale[i], axis = axis) 21 | out = tf.stack(out, axis = axis) 22 | out = self.reduce(out, axis = [2 * (i + 1) for i in range(self.loop)]) 23 | return out 24 | 25 | def get_config(self): 26 | config = super(AdaptivePooling, self).get_config() 27 | config["scale"] = self.scale 28 | return config 29 | 30 | class AdaptivePooling2D(AdaptivePooling): 31 | def __init__(self, scale, reduce = tf.reduce_mean, method = "bilinear", **kwargs): 32 | super(AdaptivePooling2D, self).__init__(scale = scale, reduce = reduce, **kwargs) 33 | self.method = method 34 | 35 | def call(self, inputs): 36 | shape = tf.shape(inputs)[1:3] 37 | w = tf.math.mod(shape, self.scale) 38 | w = tf.where(w != 0, tf.add(self.scale, -w), w) 39 | target_size = shape + w 40 | out = tf.image.resize(inputs, target_size, method = self.method) 41 | return super(AdaptivePooling2D, self).call(out) 42 | 43 | def get_config(self): 44 | config = super(AdaptivePooling2D, self).get_config() 45 | config["method"] = self.method 46 | return config 47 | 48 | class AdaptiveAveragePooling2D(AdaptivePooling2D): 49 | def __init__(self, scale, method = "bilinear", **kwargs): 50 | super(AdaptiveAveragePooling2D, self).__init__(scale = scale, reduce = tf.reduce_mean, method = method, **kwargs) 51 | 52 | class AdaptiveMaxPooling2D(AdaptivePooling2D): 53 | def __init__(self, scale, method = "bilinear", **kwargs): 54 | super(AdaptiveMaxPooling2D, self).__init__(scale = scale, reduce = tf.reduce_max, method = method, **kwargs) -------------------------------------------------------------------------------- /tfdet/core/ops/roi_extract.py: -------------------------------------------------------------------------------- 1 | import functools 2 | 3 | import tensorflow as tf 4 | 5 | @tf.function 6 | def roi2level(bbox, n_level, input_shape = (224, 224)): 7 | dtype = bbox.dtype 8 | x1, y1, x2, y2 = tf.split(bbox, 4, axis = -1) 9 | h = y2 - y1 10 | w = x2 - x1 11 | 12 | bbox_area = h * w 13 | image_area = tf.cast(input_shape[0] * input_shape[1], dtype) 14 | 15 | roi_level = tf.floor(tf.math.log((tf.sqrt(bbox_area)) / ((tf.cast(56., dtype) / tf.sqrt(image_area)) + tf.keras.backend.epsilon())) / tf.math.log(tf.cast(2., dtype))) 16 | roi_level = tf.cast(roi_level, tf.int8) 17 | roi_level = tf.clip_by_value(roi_level, 0, n_level - 1) 18 | roi_level = tf.squeeze(roi_level, axis = -1) 19 | return roi_level 20 | 21 | class RoiAlign(tf.keras.layers.Layer): 22 | def __init__(self, pool_size = 7, method = "bilinear", scatter = False, divisor = 100000, **kwargs): 23 | super(RoiAlign, self).__init__(**kwargs) 24 | self.pool_size = pool_size 25 | self.method = method 26 | self.divisor = divisor 27 | self.func = functools.partial(self.roi_align, pool_size = self.pool_size, method = self.method, divisor = self.divisor) if not scatter else functools.partial(self.roi_align_scatter, pool_size = self.pool_size, method = self.method) 28 | 29 | @staticmethod 30 | @tf.function 31 | def roi_align_scatter(feature, bbox_pred, image_shape = [1024, 1024], pool_size = 7, method = "bilinear"): 32 | if not isinstance(feature, list): 33 | feature = [feature] 34 | pool_size = [pool_size, pool_size] if not isinstance(pool_size, (tuple, list)) else pool_size 35 | 36 | bbox_shape = tf.shape(bbox_pred) 37 | batch_size = bbox_shape[0] 38 | num_proposals = bbox_shape[1] 39 | valid_flag = tf.reduce_any(tf.greater(bbox_pred, 0), axis = -1) 40 | negative_indices = tf.where(~valid_flag) 41 | 42 | roi_level = roi2level(bbox_pred, len(feature), image_shape) 43 | roi_level = tf.tensor_scatter_nd_update(roi_level, negative_indices, -tf.ones(tf.shape(negative_indices)[0], dtype = tf.int8)) 44 | 45 | x1, y1, x2, y2 = tf.split(bbox_pred, 4, axis = -1) 46 | bbox_pred = tf.concat([y1, x1, y2, x2], axis = -1) 47 | 48 | roi = tf.zeros([batch_size, num_proposals, *pool_size, tf.keras.backend.int_shape(feature[0])[-1]], dtype = feature[0].dtype) 49 | for level, x in enumerate(feature): 50 | level_flag = tf.equal(roi_level, level)#tf.logical_and(, valid_flag) 51 | level_indices = tf.where(level_flag) 52 | bbox_indices = tf.stop_gradient(tf.cast(level_indices[..., 0], tf.int32)) 53 | 54 | bbox = tf.stop_gradient(tf.gather_nd(bbox_pred, level_indices)) 55 | 56 | out = tf.image.crop_and_resize(image = x, boxes = bbox, box_indices = bbox_indices, crop_size = pool_size, method = method) 57 | out = tf.cast(out, x.dtype) 58 | 59 | roi = tf.tensor_scatter_nd_update(roi, level_indices, out) 60 | return roi 61 | 62 | @staticmethod 63 | @tf.function 64 | def roi_align(feature, bbox_pred, image_shape = [1024, 1024], pool_size = 7, method = "bilinear", divisor = 100000): 65 | if not isinstance(feature, list): 66 | feature = [feature] 67 | pool_size = [pool_size, pool_size] if not isinstance(pool_size, (tuple, list)) else pool_size 68 | 69 | bbox_shape = tf.shape(bbox_pred) 70 | batch_size = bbox_shape[0] 71 | num_proposals = bbox_shape[1] 72 | 73 | roi_level = roi2level(bbox_pred, len(feature), image_shape) 74 | 75 | x1, y1, x2, y2 = tf.split(bbox_pred, 4, axis = -1) 76 | bbox_pred = tf.concat([y1, x1, y2, x2], axis = -1) 77 | 78 | indices = [] 79 | roi = [] 80 | for level, x in enumerate(feature): 81 | level_indices = tf.where(tf.equal(roi_level, level)) 82 | level_indices = tf.cast(level_indices, tf.int32) 83 | indices.append(level_indices) 84 | 85 | bbox = tf.stop_gradient(tf.gather_nd(bbox_pred, level_indices)) 86 | bbox_indices = tf.stop_gradient(level_indices[:, 0]) 87 | 88 | out = tf.image.crop_and_resize(image = x, boxes = bbox, box_indices = bbox_indices, crop_size = pool_size, method = method) 89 | roi.append(out) 90 | roi = tf.concat(roi, axis = 0) 91 | 92 | #https://github.com/matterport/Mask_RCNN/blob/master/mrcnn/model.py 93 | indices = tf.concat(indices, axis = 0) 94 | ind_size = tf.shape(indices)[0] 95 | sort_indices = indices[:, 0] * divisor + indices[:, 1] 96 | sort_indices = tf.nn.top_k(sort_indices, k = ind_size).indices[::-1] 97 | new_indices = tf.gather(tf.range(ind_size), sort_indices) 98 | roi = tf.gather(roi, new_indices) 99 | 100 | roi = tf.reshape(roi, [batch_size, num_proposals, *pool_size, tf.shape(feature[0])[-1]]) 101 | return roi 102 | 103 | def call(self, inputs, image_shape = [1024, 1024]): 104 | feature, proposals = inputs 105 | roi = self.func(feature, proposals, image_shape = image_shape) 106 | return roi 107 | 108 | def get_config(self): 109 | config = super(RoiAlign, self).get_config() 110 | config["pool_size"] = self.pool_size 111 | config["method"] = self.method 112 | config["divisor"] = self.divisor 113 | return config -------------------------------------------------------------------------------- /tfdet/core/util/__init__.py: -------------------------------------------------------------------------------- 1 | from .log import metric2text, concat_text 2 | from .random import set_python_seed, set_random_seed, set_numpy_seed, set_tensorflow_seed, set_seed 3 | from .tf import get_batch_size, get_item, map_fn, convert_to_numpy, convert_to_pickle, convert_to_ragged_tensor, convert_to_tensor, py_func, to_categorical, pipeline, zip_pipeline, concat_pipeline, stack_pipeline, save_model, load_model, get_device, select_device, EMA 4 | from .wrapper import dict_function -------------------------------------------------------------------------------- /tfdet/core/util/log.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def metric2text(info = {}, summary = None, label = None, decimal = 4): 4 | info = {str(k):np.round(np.mean(v, axis = tuple(np.arange(1, np.ndim(v)))) if 1 < np.ndim(v) else v, decimal) for k, v in info.items()} 5 | summary = ["avg"] + [np.round(np.mean(v), decimal).astype(str) for v in info.values()] if summary is None else ["summary"] + [str(np.round(v, decimal)) for v in summary] 6 | info = {k:v.astype(str) for k, v in info.items()} 7 | label = list(range(len(list(info.values())[0]))) if label is None else label 8 | count = [max(max(max([len(str(l)) for l in label]), len("label")), len(summary[0]))] + [max(max(len(k), max([len(_v) for _v in np.reshape(v, -1)])), len(summary[1 + i])) for i, (k, v) in enumerate(info.items())] 9 | line = "+" + "+".join(["-" * cnt for cnt in count]) + "+" 10 | header = "|" + "|".join([t + " " * (c - len(t)) for t, c in zip(["label"] + list(info.keys()), count)]) + "|" 11 | content = "\n".join(["|" + "|".join([t + " " * (c - len(t)) for t, c in zip([str(l)] + [np.mean(v[i]) if 0 < np.ndim(v[i]) else v[i] for v in info.values()], count)]) + "|" for i, l in enumerate(label)]) 12 | summary = "|" + "|".join([t + " " * (c - len(t)) for t, c in zip(summary, count)]) + "|" 13 | form = "\n".join([line, 14 | header, 15 | line, 16 | content, 17 | line, 18 | summary, 19 | line]) 20 | return form 21 | 22 | def concat_text(text, axis = 1, n_blank = 1): 23 | text = [text] if np.ndim(text) == 0 else text 24 | if axis == 0: 25 | new_text = ("\n" * (n_blank + 1)).join(text) 26 | else: 27 | text = [t.split("\n") for t in text] 28 | max_row = max([len(t) for t in text]) 29 | max_cols = [max([len(t) for t in t]) for t in text] 30 | 31 | blank = " " * n_blank 32 | new_text = [] 33 | for i in range(max_row): 34 | new_line = [] 35 | for max_col, t in zip(max_cols, text): 36 | line = " " * max_col 37 | if i < len(t): 38 | line = t[i] + line[len(t[i]):] 39 | new_line.append(line) 40 | new_text.append(blank.join(new_line)) 41 | new_text = "\n".join(new_text) 42 | return new_text -------------------------------------------------------------------------------- /tfdet/core/util/random.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | 4 | import numpy as np 5 | import tensorflow as tf 6 | 7 | from tfdet.core.util.tf import get_device 8 | 9 | def set_python_seed(seed = 0): 10 | os.environ["PYTHONHASHSEED"] = str(seed) 11 | 12 | def set_random_seed(seed = 0): 13 | random.seed(seed) 14 | 15 | def set_numpy_seed(seed = 0): 16 | np.random.seed(seed) 17 | 18 | def set_tensorflow_seed(seed = 0, determinism = False): 19 | """ 20 | # This is the random seed initialization code that has to be at the beginning. 21 | """ 22 | if determinism: 23 | tf_version = float(".".join(tf.__version__.split(".")[:2])) 24 | if hasattr(tf.config.experimental, "enable_op_determinism"): 25 | tf.config.experimental.enable_op_determinism() 26 | else: 27 | os.environ["TF_DETERMINISTIC_OPS"] = "1" 28 | if tf_version == 2.9: 29 | os.environ["TF_DISABLE_DEPTHWISE_CONV_DETERMINISM_EXCEPTIONS"] = "1" 30 | if 0 < len(get_device("gpu")): 31 | os.environ["TF_CUDNN_DETERMINISTIC"] = "1" 32 | try: 33 | from tfdeterminism import patch 34 | patch() 35 | except: 36 | print("Please install 'tensorflow-determinism', and it will be more specific.") 37 | tf.random.set_seed(seed) 38 | tf.config.threading.set_inter_op_parallelism_threads(1) 39 | tf.config.threading.set_intra_op_parallelism_threads(1) 40 | 41 | def set_seed(seed = 0, determinism = False): 42 | """ 43 | # This is the random seed initialization code that has to be at the beginning. 44 | """ 45 | set_python_seed(seed) 46 | set_random_seed(seed) 47 | set_numpy_seed(seed) 48 | set_tensorflow_seed(seed, determinism = determinism) -------------------------------------------------------------------------------- /tfdet/core/util/wrapper.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import inspect 3 | 4 | import numpy as np 5 | 6 | def dict_function(function = None, keys = []): 7 | function = [function] if function is not None and not isinstance(function, (tuple, list)) else function 8 | keys = [keys] if not isinstance(keys, (tuple, list)) else keys 9 | def wrapper(function): 10 | function = [function] if np.ndim(function) == 0 else function 11 | def run(*args, **kwargs): 12 | args = list(args) 13 | map_func = args.pop(0) if 0 < len(args) and (callable(args[0]) or (isinstance(args[0], (tuple, list)) and callable(args[0][0]))) else None #[func] or func or None 14 | map_func = [map_func] if map_func is not None and not isinstance(map_func, (tuple, list)) else map_func #[func] or None 15 | return_keys = inspect.getfullargspec((map_func[0] if map_func is not None and 0 < len(map_func) else function[0])).args if len(keys) == 0 else keys 16 | if 0 < len(args) and isinstance(args[0], dict): 17 | args_keys = list(args[0].keys()) 18 | return_keys = args_keys + [key for key in return_keys if key not in args_keys] 19 | item = args[0].items() 20 | return_args = False 21 | else: 22 | item = zip(return_keys, args) 23 | return_args = True 24 | args = {k:v for k, v in item if v is not None} 25 | 26 | run_func = [functools.partial(function[0], f) for f in map_func] if map_func is not None else function 27 | base_spec = inspect.getfullargspec(function[0]) 28 | base_keys = base_spec.args + base_spec.kwonlyargs 29 | for i, func in enumerate(run_func): 30 | if callable(func): 31 | if map_func is not None: 32 | func_spec = inspect.getfullargspec(map_func[i]) 33 | func_keys = func_spec.args + func_spec.kwonlyargs 34 | func_keys = func_keys + [k for k in base_keys if k not in func_keys] 35 | else: 36 | func_spec = inspect.getfullargspec(func) 37 | func_keys = func_spec.args + func_spec.kwonlyargs 38 | func_kwargs = {k:v for k, v in kwargs.items() if k in func_keys} 39 | values = func(**args, **func_kwargs) 40 | if not isinstance(values, (tuple, list)): 41 | values = (values,) 42 | args = {k:v for k, v in zip(return_keys, values)} 43 | if return_args: 44 | result = [args[key] for key in return_keys if key in args] 45 | result = [r for r in result if r is not None] 46 | if len(result) == 0: 47 | result = None 48 | elif len(result) == 1: 49 | result = result[0] 50 | else: 51 | result = tuple(result) 52 | else: 53 | result = {k:v for k, v in args.items() if k in return_keys and v is not None} 54 | return result 55 | return run 56 | if function is not None: 57 | if 0 < len(function) and not isinstance(function[0], str): #and callable(function[0]): 58 | return wrapper(function) 59 | else: 60 | keys = function 61 | function = None 62 | return wrapper -------------------------------------------------------------------------------- /tfdet/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | from .dataset import multi_transform, Dataset, PipeLoader, GenPipeLoader, SequenceLoader 2 | from .yolo import YoloDataset 3 | 4 | from . import pipeline 5 | from . import transform 6 | from . import util 7 | 8 | from . import balloon 9 | from . import coco 10 | from . import pascal_voc -------------------------------------------------------------------------------- /tfdet/dataset/balloon.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import os 3 | 4 | import cv2 5 | import numpy as np 6 | import tensorflow as tf 7 | 8 | from .dataset import Dataset 9 | from .util.file import load_json 10 | from tfdet.core.util import pipeline, py_func 11 | 12 | LABEL = ["background", #background 13 | "balloon"] 14 | 15 | memory = {} 16 | 17 | def clear(path = None): 18 | global memory 19 | if isinstance(path, str): 20 | key = os.path.basename(path) 21 | if key in memory: 22 | del memory[key] 23 | else: 24 | memory.clear() 25 | return memory 26 | 27 | def get(path, refresh = False): 28 | global memory 29 | key = os.path.abspath(path) 30 | if refresh or key not in memory: 31 | balloon = load_json(path) 32 | balloon = list(balloon.values()) 33 | memory[key] = balloon 34 | else: 35 | balloon = memory[key] 36 | return balloon 37 | 38 | def load_info(path, x_true): 39 | balloon = get(path) 40 | data_path = os.path.dirname(os.path.abspath(path)) 41 | anno = balloon[x_true] 42 | 43 | x_true = os.path.join(data_path, anno["filename"]) 44 | y_true = [r["shape_attributes"] for r in (anno["regions"].values() if isinstance(anno["regions"], dict) else anno["regions"])] 45 | return x_true, y_true 46 | 47 | def load_annotation(x_true, y_true, mask = False): 48 | try: 49 | import skimage.draw 50 | except Exception as e: 51 | print("If you want to use 'balloon dataset', please install 'skimage'") 52 | raise e 53 | 54 | h, w = np.shape(cv2.imread(x_true, -1))[:2] 55 | new_y_true = np.array([["balloon"]] * len(y_true), dtype = np.object0) 56 | bbox_true = [] 57 | mask_true = np.zeros((len(y_true), h, w, 1), dtype = np.uint8) 58 | for i, poly in enumerate(y_true): 59 | rr, cc = skimage.draw.polygon(poly['all_points_y'], poly['all_points_x']) 60 | rr = np.clip(rr, 0, h - 1) 61 | cc = np.clip(cc, 0, w - 1) 62 | mask_true[i, rr, cc, 0] = 1 63 | pos = np.where(0 < mask_true[i, ..., 0])[:2] 64 | bbox = [np.min(pos[1]), np.min(pos[0]), np.max(pos[1]), np.max(pos[0])] 65 | bbox_true.append(bbox) 66 | y_true = new_y_true 67 | bbox_true = np.array(bbox_true, dtype = np.int32) 68 | return (x_true, y_true, bbox_true, mask_true) if mask else (x_true, y_true, bbox_true) 69 | 70 | def load_object(path, x_true, mask = False): 71 | x_true, y_true = load_info(path, x_true) 72 | return load_annotation(x_true, y_true, mask = mask) 73 | 74 | def load_dataset(path, mask = False, 75 | transform = None, refresh = False, shuffle = False, 76 | cache = None): 77 | """ 78 | https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip 79 | 80 | 81 | 1. all-in-one 82 | > dataset = tfdet.dataset.balloon.load_dataset("./balloon/train/via_region_data.json", 83 | mask = False, 84 | transform = [{"name":"load"}, 85 | {"name":"resize", "image_shape":[512, 512]}, 86 | {"name":"pad", "image_shape":[512, 512]}, 87 | {"name":"filter_annotation"}, 88 | {"name":"label_encode", "label":tfdet.dataset.balloon.LABEL}, 89 | {"name":"normalize", "mean":[123.675, 116.28, 103.53], "std":[58.395, 57.12, 57.375]}] 90 | shuffle = False, cache = "balloon_train.cache") 91 | > dataset[i] #or next(iter(dataset)) 92 | 93 | 2. split 94 | > dataset = tfdet.dataset.balloon.load_dataset("./balloon/train/via_region_data.json", 95 | mask = False, 96 | shuffle = False, cache = "balloon_train.cache") 97 | > dataset = tfdet.dataset.Dataset(dataset, 98 | transform = [{"name":"load"}, 99 | {"name":"resize", "image_shape":[512, 512]}, 100 | {"name":"pad", "image_shape":[512, 512]}, 101 | {"name":"filter_annotation"}, 102 | {"name":"label_encode", "label":tfdet.dataset.balloon.LABEL}, 103 | {"name":"normalize", "mean":[123.675, 116.28, 103.53], "std":[58.395, 57.12, 57.375]}]) 104 | > dataset[i] #or next(iter(dataset)) 105 | 106 | 3. dataset to pipe 107 | > pipe = tfdet.dataset.PipeLoader(dataset) 108 | > pipe = tfdet.dataset.pipeline.args2dict(pipe) #for train_model 109 | > pipe = tfdet.dataset.pipeline.collect(pipe) #filtered item by key 110 | > pipe = tfdet.dataset.pipeline.cast(pipe) 111 | > pipe = tfdet.dataset.pipeline.key_map(pipe, batch_size = 16, shuffle = False, prefetch = True) 112 | > next(iter(dataset)) 113 | """ 114 | if isinstance(cache, str) and os.path.exists(cache): 115 | return Dataset(transform = transform, shuffle = shuffle, cache = cache, keys = ["x_true", "y_true", "bbox_true", "mask_true"]) 116 | else: 117 | balloon = get(path, refresh = refresh) 118 | indices = np.arange(len(balloon)) 119 | object_func = functools.partial(load_object, path = path, mask = mask) 120 | return Dataset(indices, preprocess = object_func, transform = transform, shuffle = shuffle, cache = cache, keys = ["x_true", "y_true", "bbox_true", "mask_true"]) -------------------------------------------------------------------------------- /tfdet/dataset/pipeline/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | from .transform import (load, normalize, unnormalize, filter_annotation, label_encode, label_decode, compose, 3 | resize, pad, trim, crop, flip, 4 | random_crop, random_flip, multi_scale_flip, 5 | yolo_hsv, random_perspective, 6 | mosaic, mosaic9, cut_mix, cut_out, mix_up, 7 | copy_paste, remove_background, yolo_augmentation, mmdet_augmentation, 8 | key_map, collect, cast, reshape, args2dict, dict2args) 9 | try: 10 | from .transform import albumentations, weak_augmentation 11 | except: 12 | pass 13 | """ 14 | from .transform import (multi_scale_flip, key_map, collect, cast, reshape, args2dict, dict2args) 15 | from .util import pipe, zip_pipe, concat_pipe, stack_pipe, dict_py_func, dict_tf_func -------------------------------------------------------------------------------- /tfdet/dataset/pipeline/util.py: -------------------------------------------------------------------------------- 1 | import functools 2 | 3 | import numpy as np 4 | import tensorflow as tf 5 | 6 | from tfdet.core.util import (pipeline, convert_to_numpy as convert_to_numpy_eager, convert_to_pickle as convert_to_pickle_eager, py_func, dict_function, 7 | zip_pipeline as zip_pipe, concat_pipeline as concat_pipe, stack_pipeline as stack_pipe) 8 | 9 | @dict_function(keys = ["x_true", "y_true", "bbox_true", "mask_true"]) 10 | def dict_py_func(function, *args, Tout = tf.float32, **kwargs): 11 | return py_func(function, *args, Tout = Tout, **kwargs) 12 | 13 | @dict_function(keys = ["x_true", "y_true", "bbox_true", "mask_true"]) 14 | def dict_tf_func(function, *args, **kwargs): 15 | return function(*args, **kwargs) 16 | 17 | def func_format(x_true, y_true = None, bbox_true = None, mask_true = None): 18 | result = [v for v in [x_true, y_true, bbox_true, mask_true] if v is not None] 19 | result = result[0] if len(result) == 1 else tuple(result) 20 | return result 21 | 22 | def convert_to_pickle(*args): 23 | if args and isinstance(args[0], dict): 24 | out = tf.py_function(convert_to_pickle_eager, tuple(args[0].values()), Tout = [tf.string] * len(args[0])) 25 | out = {k:v for k, v in zip(args[0].keys(), out)} 26 | else: 27 | out = tf.py_function(convert_to_pickle_eager, args, Tout = [tf.string] * len(args)) 28 | out = out[0] if len(out) == 1 else tuple(out) 29 | return out 30 | 31 | def pipe(x_true, y_true = None, bbox_true = None, mask_true = None, function = None, 32 | batch_size = 0, repeat = 1, shuffle = False, prefetch = False, 33 | cache = False, num_parallel_calls = True, 34 | py_func = dict_py_func, tf_func = False, dtype = None, 35 | **kwargs): 36 | args = [arg for arg in [x_true, y_true, bbox_true, mask_true] if arg is not None] 37 | args = args[0] if len(args) == 1 else tuple(args) 38 | func = None 39 | if callable(function): 40 | if tf_func: 41 | if callable(tf_func): 42 | func = functools.partial(dict_tf_func, function, **kwargs) 43 | else: 44 | func = functools.partial(function, **kwargs) 45 | else: 46 | if dtype is None: 47 | if isinstance(x_true, tf.data.Dataset): 48 | sample_args = next(iter(x_true)) 49 | if isinstance(sample_args, dict): 50 | sample_args = convert_to_numpy_eager(sample_args) 51 | sample_result = function(**sample_args, **kwargs) 52 | else: 53 | sample_args = convert_to_numpy_eager([v for v in (sample_args if 0 < np.ndim(sample_args) else [sample_args])]) 54 | sample_result = function(*sample_args, **kwargs) 55 | else: 56 | if isinstance(x_true, dict): 57 | sample_result = function(**{k:v[0] for k, v in x_true.items()}, **kwargs) 58 | else: 59 | sample_args = [v[0] for v in [x_true, y_true, bbox_true, mask_true] if v is not None] 60 | sample_result = function(*sample_args, **kwargs) 61 | dtype = [tf.convert_to_tensor(r).dtype for r in (sample_result if isinstance(sample_result, tuple) else (sample_result,))] 62 | elif np.ndim(dtype) == 0: 63 | if isinstance(x_true, tf.data.Dataset): 64 | if isinstance(x_true.element_spec, tuple) or isinstance(x_true.element_spec, dict): 65 | dtype = [dtype] * len(x_true.element_spec) 66 | else: 67 | dtype = [dtype] * len(x_true if isinstance(x_true, dict) else [arg for arg in [x_true, y_true, bbox_true, mask_true] if arg is not None]) 68 | if 0 < np.ndim(dtype): 69 | dtype = dtype[0] if len(dtype) == 1 else tuple(dtype) 70 | func = functools.partial(py_func, function, Tout = dtype, **kwargs) if callable(function) else None 71 | return pipeline(args, function = func, 72 | batch_size = batch_size, repeat = repeat, shuffle = shuffle, prefetch = prefetch, 73 | cache = cache, num_parallel_calls = num_parallel_calls) -------------------------------------------------------------------------------- /tfdet/dataset/transform/__init__.py: -------------------------------------------------------------------------------- 1 | from .augment import random_crop, random_flip, yolo_hsv, random_perspective, mosaic, mosaic9, cut_mix, cut_out, mix_up, copy_paste, remove_background 2 | from .bundle import yolo_augmentation, mmdet_augmentation 3 | from .common import load, normalize, unnormalize, filter_annotation, label_encode, label_decode, resize, pad, trim, crop, flip, mask_downscale, compose, random_apply, random_shuffle_apply 4 | from .formatting import key_map, collect, cast, reshape, args2dict, dict2args 5 | try: 6 | from .augment import albumentations 7 | from .bundle import weak_augmentation 8 | except: 9 | pass -------------------------------------------------------------------------------- /tfdet/dataset/transform/formatting.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def key_map(x_true, y_true = None, bbox_true = None, mask_true = None, map = {"x_true":"x_true", "y_true":"y_true", "bbox_true":"bbox_true", "mask_true":"mask_true"}): 4 | if isinstance(x_true, dict): 5 | x_true = {map[k] if k in map else k:v for k, v in x_true.items()} 6 | result = [v for v in [x_true, y_true, bbox_true, mask_true] if v is not None] 7 | result = result[0] if len(result) == 1 else tuple(result) 8 | return result 9 | 10 | def collect(x_true, y_true = None, bbox_true = None, mask_true = None, keys = ["x_true", "y_true", "bbox_true", "mask_true"]): 11 | if isinstance(x_true, dict): 12 | x_true = {k:x_true[k] for k in keys if k in x_true} 13 | else: 14 | if "x_true" not in keys: 15 | x_true = None 16 | if "y_true" not in keys: 17 | y_true = None 18 | if "bbox_true" not in keys: 19 | bbox_true = None 20 | if "mask_true" not in keys: 21 | mask_true = None 22 | result = [v for v in [x_true, y_true, bbox_true, mask_true] if v is not None] 23 | result = result[0] if len(result) == 1 else tuple(result) 24 | return result 25 | 26 | def cast(x_true, y_true = None, bbox_true = None, mask_true = None, map = {"x_true":tf.float32, "y_true":tf.float32, "bbox_true":tf.float32, "mask_true":tf.float32}): 27 | if isinstance(x_true, dict): 28 | x_true = {k:tf.cast(v, map[k]) if k in map and map[k] is not None else v for k, v in x_true.items()} 29 | else: 30 | if "x_true" in map and map["x_true"] is not None: 31 | x_true = tf.cast(x_true, map["x_true"]) 32 | if y_true is not None and "y_true" in map and map["y_true"] is not None: 33 | y_true = tf.cast(y_true, map["y_true"]) 34 | if bbox_true is not None and "bbox_true" in map and map["bbox_true"] is not None: 35 | bbox_true = tf.cast(bbox_true, map["bbox_true"]) 36 | if mask_true is not None and "mask_true" in map and map["mask_true"] is not None: 37 | mask_true = tf.cast(mask_true, map["mask_true"]) 38 | result = [v for v in [x_true, y_true, bbox_true, mask_true] if v is not None] 39 | result = result[0] if len(result) == 1 else tuple(result) 40 | return result 41 | 42 | def reshape(x_true, y_true = None, bbox_true = None, mask_true = None, map = {"x_true":None, "y_true":None, "bbox_true":None, "mask_true":None}): 43 | if isinstance(x_true, dict): 44 | x_true = {k:tf.reshape(v, map[k]) if k in map and map[k] is not None else v for k, v in x_true.items()} 45 | else: 46 | if "x_true" in map and map["x_true"] is not None: 47 | x_true = tf.reshape(x_true, map["x_true"]) 48 | if y_true is not None and "y_true" in map and map["y_true"] is not None: 49 | y_true = tf.reshape(y_true, map["y_true"]) 50 | if bbox_true is not None and "bbox_true" in map and map["bbox_true"] is not None: 51 | bbox_true = tf.reshape(bbox_true, map["bbox_true"]) 52 | if mask_true is not None and "mask_true" in map and map["mask_true"] is not None: 53 | mask_true = tf.reshape(mask_true, map["mask_true"]) 54 | result = [v for v in [x_true, y_true, bbox_true, mask_true] if v is not None] 55 | result = result[0] if len(result) == 1 else tuple(result) 56 | return result 57 | 58 | def args2dict(x_true, y_true = None, bbox_true = None, mask_true = None, keys = ["x_true", "y_true", "bbox_true", "mask_true"]): 59 | if not isinstance(x_true, dict): 60 | x_true = {k:v for k, v in zip(keys, [x_true, y_true, bbox_true, mask_true]) if v is not None} 61 | y_true = bbox_true = mask_true = None 62 | result = [v for v in [x_true, y_true, bbox_true, mask_true] if v is not None] 63 | result = result[0] if len(result) == 1 else tuple(result) 64 | return result 65 | 66 | def dict2args(x_true, y_true = None, bbox_true = None, mask_true = None, keys = None): 67 | if isinstance(x_true, dict): 68 | if keys is not None: 69 | x_true = {k:x_true[k] for k in keys if k in x_true} 70 | result = list(x_true.values()) 71 | else: 72 | result = [v for v in [x_true, y_true, bbox_true, mask_true] if v is not None] 73 | result = result[0] if len(result) == 1 else tuple(result) 74 | return result -------------------------------------------------------------------------------- /tfdet/dataset/util/__init__.py: -------------------------------------------------------------------------------- 1 | from .file import list_dir, walk_dir, tree_dir, load_file, save_file, load_csv, save_csv, load_json, save_json, load_yaml, save_yaml, load_pickle, save_pickle 2 | from .image import load_image, save_image, instance2semantic, instance2bbox, instance2panoptic, panoptic2instance, trim_bbox 3 | from .numpy import pad 4 | from .xml import xml2dict, dict2xml 5 | 6 | from tfdet.core.util import convert_to_numpy, convert_to_pickle, convert_to_ragged_tensor, convert_to_tensor -------------------------------------------------------------------------------- /tfdet/dataset/util/file.py: -------------------------------------------------------------------------------- 1 | import csv 2 | import os 3 | import json 4 | import pickle 5 | import yaml 6 | 7 | import numpy as np 8 | 9 | def list_dir(path, keyword = None, absolute = False): 10 | if isinstance(keyword, str): 11 | keyword = [keyword] 12 | 13 | result = [] 14 | for file in os.listdir(path): 15 | if keyword is not None and not any([key in file for key in keyword]): 16 | continue 17 | 18 | file_path = os.path.join(path, file) 19 | if absolute: 20 | file_path = os.path.abspath(file_path) 21 | result.append(file_path) 22 | return result 23 | 24 | def walk_dir(path, keyword = None, absolute = False): 25 | if isinstance(keyword, str): 26 | keyword = [keyword] 27 | 28 | result = [] 29 | for root, dirs, files in os.walk(path): 30 | for file in files: 31 | if keyword is not None and not any([key in file for key in keyword]): 32 | continue 33 | 34 | file_path = os.path.join(root, file) 35 | if absolute: 36 | file_path = os.path.abspath(file_path) 37 | result.append(file_path) 38 | return result 39 | 40 | def tree_dir(path, file = True, keyword = None, n_skip = 4, n_blank = 4): 41 | if isinstance(keyword, str): 42 | keyword = [keyword] 43 | blank = " " * n_blank 44 | 45 | result = [] 46 | for root, dirs, files in os.walk(path): 47 | level = root.replace(path, "").count(os.sep) 48 | prefix = blank * level 49 | sub_prefix = blank * (level + 1) 50 | result.append("{0}{1}{2}".format(prefix, os.path.basename(root), os.sep)) 51 | if file: 52 | for i, f in enumerate(files): 53 | if keyword is not None and not any([key in f for key in keyword]): 54 | continue 55 | if n_skip < 1 or (i + 1) < n_skip: 56 | result.append("{0}{1}".format(sub_prefix, f)) 57 | else: 58 | result.append("{0}...".format(sub_prefix)) 59 | break 60 | return "\n".join(result) 61 | 62 | def load_file(path, map = {"\n":""}, mode = "rt"): 63 | result = [] 64 | with open(path, mode) as file: 65 | for line in file.readlines(): 66 | if isinstance(map, dict): 67 | for k, v in map.items(): 68 | line = line.replace(k, v) 69 | result.append(line) 70 | return result 71 | 72 | def save_file(data, path, end = "\n", mode = "wt"): 73 | if isinstance(data, np.ndarray): 74 | data = [data] if np.ndim(data) < 1 else data 75 | else: 76 | data = [data] if not isinstance(data, (tuple, list)) else data 77 | with open(path, mode) as file: 78 | for d in data: 79 | d = str(d) 80 | file.write("{0}{1}".format(d, end) if not (end == "" or d[-len(end):] == end) else d) 81 | return path 82 | 83 | def load_csv(path, delimiter = ",", mode = "rt"): 84 | with open(path, mode) as file: 85 | result = [line for line in csv.reader(file, delimiter = delimiter)] 86 | return result 87 | 88 | def save_csv(data, path, delimiter = ",", newline = "", mode = "wt"): 89 | if isinstance(data, np.ndarray): 90 | data = [data] if np.ndim(data) < 1 else data 91 | else: 92 | data = [data] if not isinstance(data, (tuple, list)) else data 93 | if isinstance(data[0], np.ndarray): 94 | data = [data] if np.ndim(data[0]) < 1 else data 95 | else: 96 | data = [data] if not isinstance(data[0], (tuple, list)) else data 97 | with open(path, mode, newline = newline) as file: 98 | writer = csv.writer(file, delimiter = delimiter) 99 | for row in data: 100 | writer.writerow(row) 101 | return path 102 | 103 | def load_json(path, mode = "rt"): 104 | with open(path, mode) as file: 105 | result = json.load(file) 106 | return result 107 | 108 | def save_json(data, path, mode = "wt"): 109 | with open(path, mode) as file: 110 | json.dump(data, file) 111 | return path 112 | 113 | def load_yaml(path, mode = "rt"): 114 | with open(path, mode) as file: 115 | result = yaml.full_load(file) 116 | return result 117 | 118 | def save_yaml(data, path, mode = "wt"): 119 | with open(path, mode) as file: 120 | yaml.dump(data, file) 121 | return result 122 | 123 | def load_pickle(path, mode = "rb"): 124 | with open(path, mode = mode) as file: 125 | result = pickle.loads(file.read()) 126 | return result 127 | 128 | def save_pickle(data, path , mode = "wb"): 129 | with open(path, mode = mode) as file: 130 | file.write(pickle.dumps(data)) 131 | return path 132 | -------------------------------------------------------------------------------- /tfdet/dataset/util/numpy.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | 4 | def pad(data, pad_width = None, val = 0): 5 | if pad_width is not None and 0 < np.max(pad_width): 6 | data = np.array(data) if not isinstance(data, np.ndarray) else data 7 | shape = np.shape(data) 8 | dummy = np.ndim(pad_width) < 2 9 | pad_width = [pad_width] if dummy else pad_width 10 | pad_width = list(pad_width) + [[0, 0]] * (len(shape) - len(pad_width)) 11 | new_shape = [s + sum(p) for s, p in zip(shape, pad_width)] 12 | if val is None: 13 | pad_data = np.empty(new_shape, dtype = data.dtype) 14 | else: 15 | pad_data = np.full(new_shape, val, dtype = data.dtype if np.issubdtype(type(val) if not isinstance(val, np.ndarray) else val.dtype, np.number) else np.object0) 16 | if np.all([s != 0 for s in shape]): 17 | region = tuple([slice(None if l == 0 else l, None if r == 0 else -r) for l, r in pad_width]) 18 | pad_data[region if not dummy else region[0]] = data 19 | data = pad_data 20 | return data -------------------------------------------------------------------------------- /tfdet/dataset/util/xml.py: -------------------------------------------------------------------------------- 1 | import os 2 | import xml.etree.ElementTree 3 | 4 | def xml2dict(data, attrib_prf = "@", unknown_text = "#text"): 5 | if not isinstance(data, str): 6 | raise Exception("Type incorrect") 7 | 8 | if os.path.exists(data): 9 | element = xml.etree.ElementTree.parse(data).getroot() 10 | else: 11 | element = xml.etree.ElementTree.fromstring(data) 12 | 13 | result = {} 14 | def parse(element, result): 15 | key = element.tag 16 | value = element.attrib.copy() 17 | for k in list(value.keys()): 18 | value[attrib_prf + k] = value.pop(k) 19 | 20 | childs = element.getchildren() if hasattr(element, "getchildren") else list(element) 21 | for child in childs: 22 | parse(child, value) 23 | 24 | if not childs: 25 | if not value: 26 | value = element.text 27 | elif element.text is not None: 28 | value[unknown_text] = element.text 29 | 30 | if key in result.keys(): 31 | if not isinstance(result[key], list): 32 | result[key] = [result[key]] 33 | result[key].append(value) 34 | else: 35 | result[key] = value 36 | parse(element, result) 37 | return result 38 | 39 | def dict2xml(data, save_path = None): 40 | key, val = list(data.items())[0] 41 | root = xml.etree.ElementTree.Element(key) 42 | 43 | def push(data, element): 44 | if isinstance(data, dict): 45 | for k, v in data.items(): 46 | if not isinstance(v, list): 47 | v = [v] 48 | for _v in v: 49 | push(_v, xml.etree.ElementTree.SubElement(element, k)) 50 | else: 51 | element.text = data 52 | push(val, root) 53 | 54 | if isinstance(save_path, str): 55 | tree = xml.etree.ElementTree.ElementTree(root) 56 | tree.write(save_path, "utf-8") 57 | return root 58 | -------------------------------------------------------------------------------- /tfdet/export/__init__.py: -------------------------------------------------------------------------------- 1 | from .onnx import tf2onnx, CalibrationDataset, onnx2quantize, load_onnx 2 | from .tensorrt import tf2trt, onnx2trt, add_trt_nms, load_trt, load_trt_tf, load_trt_onnx 3 | from .tf import save_tf, load_tf, tf2lite, load_tflite -------------------------------------------------------------------------------- /tfdet/export/onnx.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import tensorflow as tf 5 | 6 | def tf2onnx(model, path, opset = 13): 7 | """ 8 | - opset 9 | default = 13 10 | tflite = 16 11 | saved_model = 17 12 | """ 13 | try: 14 | import tf2onnx as onnx_converter 15 | except Exception as e: 16 | print("If you want to use 'tf2onnx', please install 'tf2onnx'") 17 | raise e 18 | 19 | name, ext = os.path.splitext(path) 20 | if len(ext) < 2: 21 | path = "{0}{1}".format(name, ".onnx") 22 | 23 | if isinstance(model, tf.keras.Model): 24 | model_proto, _ = onnx_converter.convert.from_keras(model, opset = opset, output_path = path) 25 | elif isinstance(model, str): 26 | name, ext = os.path.splitext(model) 27 | try: 28 | if "tflite" in ext: 29 | raise 30 | model = tf.keras.models.load_model(model) 31 | return tf2onnx(model, path) 32 | except: 33 | if len(ext) < 2: 34 | model = "{0}{1}".format(name, ".tflite") 35 | #model_proto, _ = onnx_converter.convert.from_tflite(model, opset = opset, output_path = path) #model output bug 36 | model_proto, _ = onnx_converter.convert._convert_common(None, name = model, opset = opset, tflite_path = model, output_path = path) 37 | else: 38 | model_proto, _ = onnx_converter.convert.from_function(model, opset = opset, output_path = path) 39 | return path 40 | 41 | class CalibrationDataset:#(onnxruntime.quantization.CalibrationDataReader): 42 | def __init__(self, data, session): 43 | if isinstance(session, list) or isinstance(session, tuple): 44 | keys = list(session) 45 | else: 46 | if isinstance(session, str): 47 | session = load_onnx(session, predict = False) 48 | keys = [inp.name for inp in session.get_inputs()] 49 | self.keys = keys 50 | del session 51 | data = data if isinstance(data, list) or isinstance(data, tuple) else [data] 52 | self.dataset = [{k:np.expand_dims(v, axis = 0) for k, v in zip(self.keys, d)} for d in zip(*data)] 53 | self.generator = iter(self.dataset) 54 | 55 | def get_next(self): 56 | return next(self.generator, None) 57 | 58 | def onnx2quantize(path, save_path, data = None): 59 | try: 60 | import onnxruntime.quantization 61 | except Exception as e: 62 | print("If you want to use 'onnx2quantize', please install 'onnxruntime'") 63 | raise e 64 | 65 | name, ext = os.path.splitext(save_path) 66 | if len(ext) < 2: 67 | save_path = "{0}{1}".format(name, ".onnx") 68 | 69 | if data is None: 70 | onnxruntime.quantization.quantize_dynamic(path, save_path) 71 | else: 72 | calibration_dataset = CalibrationDataset(data, path) if not isinstance(data, CalibrationDataset) else data 73 | onnxruntime.quantization.quantize_static(path, save_path, calibration_dataset) 74 | return save_path 75 | 76 | def load_onnx(path, gpu = None, n_thread = None, tensorrt = False, predict = True): 77 | try: 78 | import onnxruntime 79 | except Exception as e: 80 | print("If you want to use 'load_onnx', please install 'onnxruntime'") 81 | raise e 82 | 83 | name, ext = os.path.splitext(path) 84 | if len(ext) < 2: 85 | path = "{0}{1}".format(name, ".onnx") 86 | 87 | avaliable_providers = onnxruntime.get_available_providers() 88 | provider = ["CPUExecutionProvider"] 89 | if gpu is not None: 90 | if tensorrt and "TensorrtExecutionProvider" in avaliable_providers: 91 | option = {"device_id":int(gpu), "trt_fp16_enable":True} 92 | provider = [("TensorrtExecutionProvider", option)] 93 | elif "CUDAExecutionProvider" in avaliable_providers: 94 | option = {"device_id":int(gpu), "cudnn_conv_algo_search":"EXHAUSTIVE"} 95 | #option["gpu_mem_limit"] = int(gpu_mem_limit) 96 | provider = [("CUDAExecutionProvider", option)] 97 | 98 | option = None 99 | if n_thread is not None: 100 | option = onnxruntime.SessionOptions() 101 | option.inter_op_num_threads = int(n_thread) 102 | option.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL 103 | 104 | session = onnxruntime.InferenceSession(path, sess_options = option, providers = provider) 105 | if predict: 106 | input_keys = [node.name for node in session.get_inputs()] 107 | def predict(*args, **kwargs): 108 | args = {k:v for k, v in zip(input_keys[:len(args)], args)} 109 | kwargs.update(args) 110 | pred = session.run(None, kwargs) 111 | if len(pred) == 0: 112 | pred = None 113 | elif len(pred) == 1: 114 | pred = pred[0] 115 | return pred 116 | return predict 117 | else: 118 | return session -------------------------------------------------------------------------------- /tfdet/export/tf.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import tensorflow as tf 4 | 5 | from tfdet.core.util import save_model as save_tf, load_model as load_tf 6 | 7 | def tf2lite(model, path, dtype = [tf.float32], optimizations = [tf.lite.Optimize.DEFAULT], data = None): 8 | if isinstance(model, tf.keras.Model): 9 | converter = tf.lite.TFLiteConverter.from_keras_model(model) 10 | elif isinstance(model, str): 11 | converter = tf.lite.TFLiteConverter.from_saved_model(model) 12 | else: 13 | converter = tf.lite.TFLiteConverter.from_concrete_functions(model) 14 | 15 | if optimizations is not None: 16 | if not isinstance(optimizations, list): 17 | optimizations = [optimizations] 18 | converter.optimizations = optimizations 19 | if data is None: 20 | if dtype is not None: 21 | converter.target_spec.supported_types = [dtype] if not isinstance(dtype, list) else dtype 22 | converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS] 23 | else: 24 | converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8, tf.lite.OpsSet.SELECT_TF_OPS] 25 | converter.inference_input_type = tf.int8 #tf.uint8 26 | converter.inference_output_type = tf.int8 #tf.uint8 27 | data = data if isinstance(data, list) or isinstance(data, tuple) else [data] 28 | def representative_dataset(): 29 | for ds in zip(*data): 30 | yield [np.expand_dims(d, axis = 0) for d in ds] 31 | converter.representative_dataset = representative_dataset 32 | tflite_model = converter.convert() 33 | 34 | name, ext = os.path.splitext(path) 35 | if len(ext) < 2: 36 | path = "{0}{1}".format(name, ".tflite") 37 | with open(path, "wb") as file: 38 | file.write(tflite_model) 39 | return path 40 | 41 | def load_tflite(path, n_thread = None, content = None, predict = True): 42 | name, ext = os.path.splitext(path) 43 | if len(ext) < 2: 44 | path = "{0}{1}".format(name, ".tflite") 45 | interpreter = tf.lite.Interpreter(path, model_content = content, num_threads = n_thread) 46 | if predict: 47 | interpreter.allocate_tensors() 48 | info = interpreter.get_signature_list()["serving_default"] 49 | input_keys, output_keys = info["inputs"], info["outputs"] 50 | input_details, output_details = interpreter.get_input_details(), interpreter.get_output_details() 51 | def predict(*args, **kwargs): 52 | args = {k:v for k, v in zip(input_keys[:len(args)], args)} 53 | kwargs.update(args) 54 | #kwargs = {k:v if tf.is_tensor(v) else tf.convert_to_tensor(v) for k, v in kwargs.items()} 55 | for key, detail in zip(input_keys, input_details): 56 | interpreter.set_tensor(detail["index"], kwargs[key]) 57 | interpreter.invoke() 58 | pred = [interpreter.get_tensor(detail["index"]) for key, detail in zip(output_keys, output_details)] 59 | if len(pred) == 0: 60 | pred = None 61 | elif len(pred) == 1: 62 | pred = pred[0] 63 | return pred 64 | return predict 65 | else: 66 | return interpreter -------------------------------------------------------------------------------- /tfdet/model/__init__.py: -------------------------------------------------------------------------------- 1 | from . import backbone 2 | from . import detector 3 | from . import head 4 | from . import neck 5 | from . import postprocess 6 | from . import train -------------------------------------------------------------------------------- /tfdet/model/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | from .convnext import convnext_tiny, convnext_small, convnext_base, convnext_large, convnext_xlarge 2 | from .darknet import darknet53, darknet19, csp_darknet53, csp_darknet19 3 | from .densenet import densenet121, densenet169, densenet201 4 | from .effnet import effnet_b0, effnet_b1, effnet_b2, effnet_b3, effnet_b4, effnet_b5, effnet_b6, effnet_b7 5 | from .effnet import effnet_lite_b0, effnet_lite_b1, effnet_lite_b2, effnet_lite_b3, effnet_lite_b4 6 | from .effnet import effnet_v2_b0, effnet_v2_b1, effnet_v2_b2, effnet_v2_b3, effnet_v2_s, effnet_v2_m, effnet_v2_l 7 | from .mobilenet import mobilenet_v2, mobilenet_v3_small, mobilenet_v3_large 8 | from .resnet import resnet18, resnet34, resnet50, resnet101, resnet152, resnext50_32x4d, resnext101_32x8d, resnext101_64x4d, wide_resnet50_2, wide_resnet101_2 9 | from .resnest import resnest50, resnest101, resnest200, resnest269 10 | from .vgg import vgg11, vgg11_bn, vgg13, vgg13_bn, vgg16, vgg16_bn, vgg19, vgg19_bn 11 | #from .swin_transformer import (swin_transformer_tiny_224_w7_1k, swin_transformer_tiny_224_w7_22k, swin_transformer_tiny_224_w7_22kto1k, 12 | # swin_transformer_small_224_w7_1k, swin_transformer_small_224_w7_22k, swin_transformer_small_224_w7_22kto1k, 13 | # swin_transformer_base_224_w7_1k, swin_transformer_base_224_w7_22k, swin_transformer_base_224_w7_22kto1k, 14 | # swin_transformer_base_384_w12_1k, swin_transformer_base_384_w12_22k, swin_transformer_base_384_w12_22kto1k, 15 | # swin_transformer_large_224_w7_22k, swin_transformer_large_224_w7_22kto1k, 16 | # swin_transformer_large_384_w12_22k, swin_transformer_large_384_w12_22kto1k) 17 | #from .swin_transformer import (swin_transformer_v2_tiny_256_w8_1k, swin_transformer_v2_tiny_256_w16_1k, 18 | # swin_transformer_v2_small_256_w8_1k, swin_transformer_v2_small_256_w16_1k, 19 | # swin_transformer_v2_base_256_w8_1k, swin_transformer_v2_base_256_w16_1k, 20 | # swin_transformer_v2_base_192_w12_22k, swin_transformer_v2_base_256_w16_22kto1k, swin_transformer_v2_base_384_w24_22kto1k, 21 | # swin_transformer_v2_large_192_w12_22k, swin_transformer_v2_large_256_w16_22kto1k, swin_transformer_v2_large_384_w24_22kto1k) 22 | from .swin_transformer import swin_tiny, swin_small, swin_base, swin_v2_tiny, swin_v2_small, swin_v2_base -------------------------------------------------------------------------------- /tfdet/model/detector/__init__.py: -------------------------------------------------------------------------------- 1 | from .effdet import effdet, effdet_d0, effdet_d1, effdet_d2, effdet_d3, effdet_d4, effdet_d5, effdet_d6, effdet_d7, effdet_d7x 2 | from .effdet import effdet_lite, effdet_lite_d0, effdet_lite_d1, effdet_lite_d2, effdet_lite_d3, effdet_lite_d3x, effdet_lite_d4 3 | from .fcos import fcos 4 | from .rcnn import rcnn, faster_rcnn, mask_rcnn, cascade_rcnn, hybrid_task_cascade_rcnn 5 | from .retina import retinanet 6 | from .yolo import yolo_v3, yolo_tiny_v3, yolo_v4, yolo_tiny_v4 7 | hybrid_task_cascade = hybrid_task_cascade_rcnn 8 | 9 | from .spade import spade 10 | from .padim import padim 11 | from .patch_core import patch_core 12 | 13 | from .deeplab import deeplab_v3, deeplab_v3_plus 14 | from .fcn import fcn, aux_fcn 15 | from .pspnet import pspnet 16 | from .unet import unet, unet_2plus 17 | from .upernet import upernet -------------------------------------------------------------------------------- /tfdet/model/detector/deeplab.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..head import deeplab_v3_head 4 | 5 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 6 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 7 | 8 | def head_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def normalize(axis = -1, momentum = 0.9, epsilon = 1e-5, **kwargs): 12 | return tf.keras.layers.BatchNormalization(axis = axis, momentum = momentum, epsilon = epsilon, **kwargs) 13 | 14 | def deeplab_v3(feature, n_class = 35, rate = [1, 12, 24, 36], n_feature = 512, method = "bilinear", dropout_rate = 0.1, 15 | neck = None, 16 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 17 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 18 | if neck is not None: 19 | feature = neck(name = "neck")(feature) 20 | if isinstance(feature, list): 21 | feature = feature[-1] 22 | out = deeplab_v3_head(feature, n_class, rate, n_feature, method = method, dropout_rate = dropout_rate, 23 | logits_convolution = logits_convolution, logits_activation = logits_activation, 24 | convolution = convolution, normalize = normalize, activation = activation) 25 | return out 26 | 27 | def deeplab_v3_plus(feature, n_class = 35, rate = [1, 12, 24, 36], n_feature = 512, n_low_feature = 48, method = "bilinear", dropout_rate = 0.1, 28 | neck = None, 29 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 30 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 31 | if neck is not None: 32 | feature = neck(name = "neck")(feature) 33 | if not isinstance(feature, list): 34 | feature = [feature] 35 | out = deeplab_v3_head(feature, n_class, rate, n_feature, n_low_feature, method = method, dropout_rate = dropout_rate, 36 | logits_convolution = logits_convolution, logits_activation = logits_activation, 37 | convolution = convolution, normalize = normalize, activation = activation) 38 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/fcn.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..head import fcn_head 4 | 5 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 6 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 7 | 8 | def head_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def normalize(axis = -1, momentum = 0.9, epsilon = 1e-5, **kwargs): 12 | return tf.keras.layers.BatchNormalization(axis = axis, momentum = momentum, epsilon = epsilon, **kwargs) 13 | 14 | def fcn(feature, n_class = 35, n_feature = 512, n_depth = 2, dropout_rate = 0.1, 15 | neck = None, 16 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 17 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 18 | if neck is not None: 19 | feature = neck(name = "neck")(feature) 20 | out = fcn_head(feature, n_class = n_class, n_feature = n_feature, n_depth = n_depth, dropout_rate = dropout_rate, 21 | logits_convolution = logits_convolution, logits_activation = logits_activation, 22 | convolution = convolution, normalize = normalize, activation = activation) 23 | return out 24 | 25 | def aux_fcn(feature, n_class = 35, n_feature = 256, n_depth = 1, dropout_rate = 0.1, 26 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 27 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 28 | if isinstance(feature, (tuple, list)): 29 | feature = feature[-2] 30 | out = fcn_head(feature, n_class = n_class, n_feature = n_feature, n_depth = n_depth, dropout_rate = dropout_rate, 31 | logits_convolution = logits_convolution, logits_activation = logits_activation, 32 | convolution = convolution, normalize = normalize, activation = activation, 33 | prefix = "aux") 34 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/fcos.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from tfdet.core.ops.initializer import PriorProbability 5 | from ..head import fcos_head 6 | from ..neck import FeatureAlign, fpn 7 | 8 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def neck_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "glorot_uniform", **kwargs): 12 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 13 | 14 | def cls_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), bias_initializer = PriorProbability(probability = 0.01), **kwargs): 15 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 16 | 17 | def bbox_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 18 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 19 | 20 | def conf_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 21 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 22 | 23 | def group_normalize(groups = 32, epsilon = 1e-5, momentum = 0.9, **kwargs): 24 | layer = None 25 | try: 26 | layer = tf.keras.layers.GroupNormalization(groups = groups, epsilon = epsilon, **kwargs) 27 | except: 28 | try: 29 | import tensorflow_addons as tfa 30 | layer = tfa.layers.GroupNormalization(groups = groups, epsilon = epsilon, **kwargs) 31 | except: 32 | pass 33 | if layer is None: 34 | print("If you want to use 'GroupNormalization', please install 'tensorflow 2.11▲ or tensorflow_addons'") 35 | layer = tf.keras.layers.BatchNormalization(momentum = momentum, epsilon = epsilon, **kwargs) 36 | return layer 37 | 38 | def neck(n_feature = 256, n_sampling = 1, pre_sampling = True, neck = fpn, neck_n_depth = 1, use_bias = None, convolution = neck_conv, normalize = None, **kwargs): 39 | return FeatureAlign(n_feature = n_feature, n_sampling = n_sampling, pre_sampling = pre_sampling, neck = neck, neck_n_depth = neck_n_depth, use_bias = use_bias, convolution = convolution, normalize = normalize, **kwargs) 40 | 41 | def fcos(feature, n_class = 21, image_shape = [1024, 1024], n_feature = 256, n_depth = 4, use_bias = None, 42 | centerness = True, 43 | neck = neck, 44 | cls_convolution = cls_conv, cls_activation = tf.keras.activations.sigmoid, 45 | bbox_convolution = bbox_conv, bbox_activation = tf.keras.activations.linear, 46 | conf_convolution = conf_conv, conf_normalize = None, conf_activation = tf.keras.activations.sigmoid, 47 | convolution = conv, normalize = group_normalize, activation = tf.keras.activations.relu): 48 | if neck is not None: 49 | feature = neck(name = "neck")(feature) 50 | out = fcos_head(feature, n_class = n_class, image_shape = image_shape, n_feature = n_feature, n_depth = n_depth, use_bias = use_bias, 51 | centerness = centerness, 52 | cls_convolution = cls_convolution, cls_activation = cls_activation, 53 | bbox_convolution = bbox_convolution, bbox_activation = bbox_activation, 54 | conf_convolution = conf_convolution, conf_normalize = conf_normalize, conf_activation = conf_activation, 55 | convolution = convolution, normalize = normalize, activation = activation) 56 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/padim.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.ops import mahalanobis 4 | from ..head import padim_head 5 | 6 | def padim(feature, mean = None, cvar_inv = None, image_shape = [224, 224], sampling_index = None, sigma = 4, metric = mahalanobis, method = "bilinear", memory_reduce = False, batch_size = 1): 7 | out = padim_head(feature, mean, cvar_inv, image_shape = image_shape, sampling_index = sampling_index, sigma = sigma, metric = metric, method = method, memory_reduce = memory_reduce, batch_size = batch_size) 8 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/patch_core.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..head import patch_core_head 4 | 5 | def patch_core(feature, feature_vector = None, image_shape = [224, 224], k = 9, sampling_index = None, pool_size = 3, sigma = 4, method = "bilinear", memory_reduce = False): 6 | out = patch_core_head(feature, feature_vector, image_shape = image_shape, k = k, sampling_index = sampling_index, pool_size = pool_size, sigma = sigma, method = method, memory_reduce = memory_reduce) 7 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/pspnet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..head import pspnet_head 4 | 5 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 6 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 7 | 8 | def head_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def normalize(axis = -1, momentum = 0.9, epsilon = 1e-5, **kwargs): 12 | return tf.keras.layers.BatchNormalization(axis = axis, momentum = momentum, epsilon = epsilon, **kwargs) 13 | 14 | def pspnet(feature, n_class = 35, n_feature = 512, pool_scale = [1, 2, 3, 6], max_pooling = False, method = "bilinear", dropout_rate = 0.1, 15 | neck = None, 16 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 17 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 18 | if neck is not None: 19 | feature = neck(name = "neck")(feature) 20 | out = pspnet_head(feature, n_class = n_class, n_feature = n_feature, pool_scale = pool_scale, max_pooling = max_pooling, method = method, dropout_rate = dropout_rate, 21 | logits_convolution = logits_convolution, logits_activation = logits_activation, 22 | convolution = convolution, normalize = normalize, activation = activation) 23 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/retina.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from tfdet.core.ops.initializer import PriorProbability 5 | from ..head import retina_head 6 | from ..neck import FeatureAlign, fpn 7 | 8 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def neck_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "glorot_uniform", **kwargs): 12 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 13 | 14 | def cls_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), bias_initializer = PriorProbability(probability = 0.01), **kwargs): 15 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 16 | 17 | def bbox_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 18 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 19 | 20 | def neck(n_feature = 256, n_sampling = 1, pre_sampling = True, neck = fpn, neck_n_depth = 1, use_bias = None, convolution = neck_conv, normalize = None, **kwargs): 21 | return FeatureAlign(n_feature = n_feature, n_sampling = n_sampling, pre_sampling = pre_sampling, neck = neck, neck_n_depth = neck_n_depth, use_bias = use_bias, convolution = convolution, normalize = normalize, **kwargs) 22 | 23 | def retinanet(feature, n_class = 21, image_shape = [1024, 1024], n_feature = 256, n_depth = 4, use_bias = None, 24 | scale = [32, 64, 128, 256, 512], ratio = [0.5, 1, 2], octave = 3, 25 | neck = neck, 26 | cls_convolution = cls_conv, cls_activation = tf.keras.activations.sigmoid, 27 | bbox_convolution = bbox_conv, bbox_activation = tf.keras.activations.linear, 28 | convolution = conv, normalize = None, activation = tf.keras.activations.relu): 29 | if neck is not None: 30 | feature = neck(name = "neck")(feature) 31 | out = retina_head(feature, n_class = n_class, image_shape = image_shape, n_feature = n_feature, n_depth = n_depth, use_bias = use_bias, 32 | scale = scale, ratio = ratio, octave = octave, 33 | cls_convolution = cls_convolution, cls_activation = cls_activation, 34 | bbox_convolution = bbox_convolution, bbox_activation = bbox_activation, 35 | convolution = convolution, normalize = normalize, activation = activation) 36 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/spade.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..head import spade_head 4 | 5 | def spade(feature, feature_vector = None, image_shape = [224, 224], k = 50, sampling_index = None, sigma = 4, method = "bilinear"): 6 | out = spade_head(feature, feature_vector, image_shape = image_shape, k = k, sampling_index = sampling_index, sigma = sigma, method = method) 7 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/unet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..head import unet_head, unet_2plus_head 4 | 5 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 6 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 7 | 8 | def head_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def normalize(axis = -1, momentum = 0.9, epsilon = 1e-5, **kwargs): 12 | return tf.keras.layers.BatchNormalization(axis = axis, momentum = momentum, epsilon = epsilon, **kwargs) 13 | 14 | def unet(feature, n_class = 35, n_feature = 64, n_depth = 5, method = "bilinear", dropout_rate = 0.1, 15 | neck = None, 16 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 17 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 18 | if neck is not None: 19 | feature = neck(name = "neck")(feature) 20 | out = unet_head(feature, n_class = n_class, n_feature = n_feature, n_depth = n_depth, method = method, dropout_rate = dropout_rate, 21 | logits_convolution = logits_convolution, logits_activation = logits_activation, 22 | convolution = convolution, normalize = normalize, activation = activation) 23 | return out 24 | 25 | def unet_2plus(feature, n_class = 35, n_feature = 512, deep_supervision = False, method = "bilinear", dropout_rate = 0.1, 26 | neck = None, 27 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 28 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 29 | if neck is not None: 30 | feature = neck(name = "neck")(feature) 31 | out = unet_2plus_head(feature, n_class = n_class, n_feature = n_feature, deep_supervision = deep_supervision, method = method, dropout_rate = dropout_rate, 32 | logits_convolution = logits_convolution, logits_activation = logits_activation, 33 | convolution = convolution, normalize = normalize, activation = activation) 34 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/upernet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..head import upernet_head 4 | 5 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 6 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 7 | 8 | def head_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def normalize(axis = -1, momentum = 0.9, epsilon = 1e-5, **kwargs): 12 | return tf.keras.layers.BatchNormalization(axis = axis, momentum = momentum, epsilon = epsilon, **kwargs) 13 | 14 | def upernet(feature, n_class = 35, n_feature = 512, pool_scale = [1, 2, 3, 6], max_pooling = False, method = "bilinear", dropout_rate = 0.1, 15 | neck = None, 16 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 17 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 18 | if neck is not None: 19 | feature = neck(name = "neck")(feature) 20 | out = upernet_head(feature, n_class = n_class, n_feature = n_feature, pool_scale = pool_scale, max_pooling = max_pooling, method = method, dropout_rate = dropout_rate, 21 | logits_convolution = logits_convolution, logits_activation = logits_activation, 22 | convolution = convolution, normalize = normalize, activation = activation) 23 | return out -------------------------------------------------------------------------------- /tfdet/model/detector/yolo.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tensorflow as tf 3 | import numpy as np 4 | 5 | from ..backbone.darknet import darknet53, darknet19, csp_darknet53, csp_darknet19, load_weight 6 | from ..head import yolo_v3_head, yolo_tiny_v3_head, yolo_v4_head, yolo_tiny_v4_head 7 | 8 | def mish(x): 9 | return x * tf.math.tanh(tf.math.softplus(x)) 10 | 11 | def leaky_relu(x, alpha = 0.1): 12 | return tf.nn.leaky_relu(x, alpha = alpha) 13 | 14 | yolo_urls = {"yolo_v3":"https://pjreddie.com/media/files/yolov3.weights", 15 | "yolo_tiny_v3":"https://pjreddie.com/media/files/yolov3-tiny.weights", 16 | "yolo_v4":"https://github.com/AlexeyAB/darknet/releases/download/yolov4/yolov4.weights", 17 | "yolo_tiny_v4":"https://github.com/AlexeyAB/darknet/releases/download/yolov4/yolov4-tiny.weights"} 18 | 19 | def yolo_v3(x, n_class = 80, size = [[ 10, 13], [ 16, 30], [ 33, 23], 20 | [ 30, 61], [ 62, 45], [ 59, 119], 21 | [116, 90], [156, 198], [373, 326]], 22 | feature_share = True, method = "nearest", 23 | normalize = tf.keras.layers.BatchNormalization, activation = mish, post_activation = leaky_relu, weights = "darknet"): 24 | """ 25 | darknet, yolo > normalize(x, rescale = 1 / 255, mean = None, std = None) 26 | """ 27 | feature = darknet53(x, csp = False, normalize = normalize, activation = activation, post_activation = post_activation, weights = None) 28 | score_pred, logit_pred, bbox_pred, anchors = yolo_v3_head(feature, n_class = n_class, image_shape = tf.shape(x)[1:3], size = size, feature_share = feature_share, method = method, normalize = normalize, activation = activation, post_activation = post_activation) 29 | 30 | if weights is not None: 31 | model = tf.keras.Model(x, [score_pred, logit_pred, bbox_pred]) 32 | if weights in ["darknet", "yolo"]: 33 | load_weight(model, yolo_urls["yolo_v3"], n_class = n_class, only_darknet = weights == "darknet") 34 | else: 35 | model.load_weights(weights) 36 | return score_pred, logit_pred, bbox_pred, anchors 37 | 38 | def yolo_tiny_v3(x, n_class = 80, size = [[23, 27], [ 37, 58], [ 81, 82], 39 | [81, 82], [135, 169], [344, 319]], 40 | feature_share = True, method = "nearest", 41 | normalize = tf.keras.layers.BatchNormalization, activation = mish, post_activation = leaky_relu, weights = "darknet"): 42 | """ 43 | darknet, yolo > normalize(x, rescale = 1 / 255, mean = None, std = None) 44 | """ 45 | feature = darknet19(x, csp = False, normalize = normalize, activation = activation, weights = None) 46 | score_pred, logit_pred, bbox_pred, anchors = yolo_tiny_v3_head(feature, n_class = n_class, image_shape = tf.shape(x)[1:3], size = size, feature_share = feature_share, method = method, normalize = normalize, activation = activation, post_activation = post_activation) 47 | 48 | if weights is not None: 49 | model = tf.keras.Model(x, [score_pred, logit_pred, bbox_pred]) 50 | if weights in ["darknet", "yolo"]: 51 | load_weight(model, yolo_urls["yolo_tiny_v3"], n_class = n_class, only_darknet = weights == "darknet") 52 | else: 53 | model.load_weights(weights) 54 | return score_pred, logit_pred, bbox_pred, anchors 55 | 56 | def yolo_v4(x, n_class = 80, size = [[ 10, 13], [ 16, 30], [ 33, 23], 57 | [ 30, 61], [ 62, 45], [ 59, 119], 58 | [116, 90], [156, 198], [373, 326]], 59 | feature_share = True, method = "nearest", 60 | normalize = tf.keras.layers.BatchNormalization, activation = mish, post_activation = leaky_relu, weights = "darknet"): 61 | """ 62 | darknet, yolo > normalize(x, rescale = 1 / 255, mean = None, std = None) 63 | """ 64 | feature = darknet53(x, csp = True, normalize = normalize, activation = activation, post_activation = post_activation, weights = None) 65 | score_pred, logit_pred, bbox_pred, anchors = yolo_v4_head(feature, n_class = n_class, image_shape = tf.shape(x)[1:3], size = size, feature_share = feature_share, method = method, normalize = normalize, activation = activation, post_activation = post_activation) 66 | 67 | if weights is not None: 68 | model = tf.keras.Model(x, [score_pred, logit_pred, bbox_pred]) 69 | if weights in ["darknet", "yolo"]: 70 | load_weight(model, yolo_urls["yolo_v4"], n_class = n_class, only_darknet = weights == "darknet") 71 | else: 72 | model.load_weights(weights) 73 | return score_pred, logit_pred, bbox_pred, anchors 74 | 75 | def yolo_tiny_v4(x, n_class = 80, size = [[23, 27], [ 37, 58], [ 81, 82], 76 | [81, 82], [135, 169], [344, 319]], 77 | feature_share = True, method = "nearest", 78 | normalize = tf.keras.layers.BatchNormalization, activation = mish, post_activation = leaky_relu, weights = "darknet"): 79 | """ 80 | darknet, yolo > normalize(x, rescale = 1 / 255, mean = None, std = None) 81 | """ 82 | feature = darknet19(x, csp = True, normalize = normalize, activation = activation, weights = None) 83 | score_pred, logit_pred, bbox_pred, anchors = yolo_tiny_v4_head(feature, n_class = n_class, image_shape = tf.shape(x)[1:3], size = size, feature_share = feature_share, method = method, normalize = normalize, activation = activation, post_activation = post_activation) 84 | 85 | if weights is not None: 86 | model = tf.keras.Model(x, [score_pred, logit_pred, bbox_pred]) 87 | if weights in ["darknet", "yolo"]: 88 | load_weight(model, yolo_urls["yolo_tiny_v4"], n_class = n_class, only_darknet = weights == "darknet") 89 | else: 90 | model.load_weights(weights) 91 | return score_pred, logit_pred, bbox_pred, anchors -------------------------------------------------------------------------------- /tfdet/model/head/__init__.py: -------------------------------------------------------------------------------- 1 | from .fcos import fcos_head 2 | from .rcnn import rpn_head, bbox_head, mask_head, semantic_head, Rpn2Proposal, Classifier2Proposal 3 | from .retina import retina_head 4 | from .yolo import yolo_v3_head, yolo_tiny_v3_head, yolo_v4_head, yolo_tiny_v4_head 5 | 6 | from .spade import spade_head 7 | from .padim import padim_head 8 | from .patch_core import patch_core_head 9 | 10 | from .deeplab import deeplab_v3_head 11 | from .fcn import fcn_head 12 | from .pspnet import pspnet_head 13 | from .unet import unet_head, unet_2plus_head 14 | from .upernet import upernet_head -------------------------------------------------------------------------------- /tfdet/model/head/fcn.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 4 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 5 | 6 | def head_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 7 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 8 | 9 | def normalize(axis = -1, momentum = 0.9, epsilon = 1e-5, **kwargs): 10 | return tf.keras.layers.BatchNormalization(axis = axis, momentum = momentum, epsilon = epsilon, **kwargs) 11 | 12 | def fcn_head(feature, n_class = 35, n_feature = 512, n_depth = 2, dropout_rate = 0.1, 13 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 14 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu, 15 | prefix = ""): 16 | #https://arxiv.org/pdf/1411.4038.pdf 17 | if isinstance(feature, (tuple, list)): 18 | feature = feature[-1] 19 | prefix = "{0}_".format(prefix) if len(prefix) != 0 else prefix 20 | 21 | out = feature 22 | for index in range(n_depth): 23 | out = convolution(n_feature, 3, padding = "same", use_bias = normalize is None, name = "{0}feature_conv{1}".format(prefix, index + 1))(out) 24 | if normalize is not None: 25 | out = normalize(name = "{0}feature_norm{1}".format(prefix, index + 1))(out) 26 | out = tf.keras.layers.Activation(activation, name = "{0}feature_act{1}".format(prefix, index + 1))(out) 27 | 28 | if 0 < n_depth: 29 | out = tf.keras.layers.Concatenate(axis = -1, name = "{0}post_concat".format(prefix))([out, feature]) 30 | out = convolution(n_feature, 3, padding = "same", use_bias = normalize is None, name = "{0}post_conv".format(prefix))(out) 31 | if normalize is not None: 32 | out = normalize(name = "{0}post_norm".format(prefix))(out) 33 | out = tf.keras.layers.Activation(activation, name = "{0}post_act".format(prefix))(out) 34 | 35 | if 0 < dropout_rate: 36 | out = tf.keras.layers.Dropout(dropout_rate)(out) 37 | out = logits_convolution(n_class, 1, use_bias = True, name = "{0}logits".format(prefix))(out) 38 | out = tf.keras.layers.Activation(logits_activation if logits_activation is not None else tf.keras.activations.linear, dtype = tf.float32, name = "{0}logits_act".format(prefix))(out) 39 | return out -------------------------------------------------------------------------------- /tfdet/model/head/padim.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import cv2 3 | import numpy as np 4 | 5 | from tfdet.core.ops import feature_extract, mahalanobis 6 | 7 | class FeatureExtractor(tf.keras.layers.Layer): 8 | def __init__(self, sampling_index = None, memory_reduce = False, **kwargs): 9 | super(FeatureExtractor, self).__init__(**kwargs) 10 | self.sampling_index = sampling_index 11 | self.memory_reduce = memory_reduce 12 | 13 | def call(self, inputs): 14 | out = feature_extract(inputs, sampling_index = self.sampling_index, pool_size = 1, sub_sampling = False, concat = True, memory_reduce = self.memory_reduce) 15 | return out 16 | 17 | def get_config(self): 18 | config = super(FeatureExtractor, self).get_config() 19 | config["sampling_index"] = self.sampling_index 20 | config["memory_reduce"] = self.memory_reduce 21 | return config 22 | 23 | class Head(tf.keras.layers.Layer): 24 | def __init__(self, mean, cvar_inv = None, image_shape = [224, 224], sigma = 4, metric = mahalanobis, method = "bilinear", batch_size = 1, **kwargs): 25 | super(Head, self).__init__(**kwargs) 26 | if isinstance(mean, (tuple, list)): 27 | mean, cvar_inv = mean 28 | self.mean = tf.constant(mean, dtype = tf.float32) if not tf.is_tensor(mean) else mean 29 | self.cvar_inv = tf.constant(cvar_inv, dtype = tf.float32) if not tf.is_tensor(cvar_inv) else cvar_inv 30 | self.image_shape = image_shape 31 | self.sigma = sigma 32 | self.metric = metric 33 | self.method = method 34 | self.batch_size = batch_size 35 | 36 | self.kernel = (2 * round(4 * sigma) + 1,) * 2 37 | 38 | def call(self, inputs): 39 | b = tf.shape(inputs)[0] 40 | h, w, c = tf.keras.backend.int_shape(inputs)[1:] 41 | feature = tf.reshape(inputs, [b, h * w, c]) 42 | mask = tf.zeros((h * w, b, 1)) 43 | mask = tf.while_loop(lambda index, mask: index < (h * w), 44 | lambda index, mask: (index + 1, tf.tensor_scatter_nd_update(mask, tf.stack([tf.ones(b, dtype = tf.int32) * index, tf.range(b)], axis = -1), mahalanobis(feature[:, index], self.mean[index], self.cvar_inv[index]))), 45 | (0, mask), 46 | parallel_iterations = self.batch_size)[1] 47 | mask = tf.reshape(tf.transpose(mask, [1, 0, 2]), [b, h, w, 1]) 48 | 49 | #upsampling 50 | mask = tf.image.resize(mask, self.image_shape, method = self.method) 51 | 52 | #gaussian smoothing 53 | if 0 < self.sigma: 54 | mask = tf.py_function(lambda *args: np.array([cv2.GaussianBlur(m, self.kernel, self.sigma) for m in args[0].numpy()]), inp = [mask], Tout = mask.dtype) 55 | mask = tf.reshape(mask, [-1, *self.image_shape, 1]) 56 | return tf.expand_dims(tf.reduce_max(mask, axis = (1, 2, 3)), axis = -1), mask 57 | 58 | def get_config(self): 59 | config = super(Head, self).get_config() 60 | config["image_shape"] = self.image_shape 61 | config["sigma"] = self.sigma 62 | config["method"] = self.method 63 | config["batch_size"] = self.batch_size 64 | return config 65 | 66 | def padim_head(feature, mean = None, cvar_inv = None, image_shape = [224, 224], sampling_index = None, sigma = 4, metric = mahalanobis, method = "bilinear", memory_reduce = False, batch_size = 1): 67 | feature = FeatureExtractor(sampling_index = sampling_index, memory_reduce = memory_reduce, name = "feature_extractor")(feature) 68 | if mean is not None: 69 | score, mask = Head(mean = mean, cvar_inv = cvar_inv, image_shape = image_shape, sigma = sigma, metric = metric, method = method, batch_size = batch_size, name = "padim")(feature) 70 | return score, mask 71 | else: 72 | return feature -------------------------------------------------------------------------------- /tfdet/model/head/patch_core.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import cv2 3 | import numpy as np 4 | 5 | from tfdet.core.ops import feature_extract, euclidean_matrix 6 | 7 | class FeatureExtractor(tf.keras.layers.Layer): 8 | def __init__(self, sampling_index = None, pool_size = 3, memory_reduce = False, **kwargs): 9 | super(FeatureExtractor, self).__init__(**kwargs) 10 | self.sampling_index = sampling_index 11 | self.pool_size = pool_size 12 | self.memory_reduce = memory_reduce 13 | 14 | def call(self, inputs): 15 | out = feature_extract(inputs, sampling_index = self.sampling_index, pool_size = self.pool_size, sub_sampling = False, concat = True, memory_reduce = self.memory_reduce) 16 | return out 17 | 18 | def get_config(self): 19 | config = super(FeatureExtractor, self).get_config() 20 | config["sampling_index"] = self.sampling_index 21 | config["pool_size"] = self.pool_size 22 | config["memory_reduce"] = self.memory_reduce 23 | return config 24 | 25 | class Head(tf.keras.layers.Layer): 26 | def __init__(self, feature_vector, image_shape = [224, 224], k = 9, sigma = 4, method = "bilinear", **kwargs): 27 | super(Head, self).__init__(**kwargs) 28 | self.feature_vector = feature_vector 29 | self.image_shape = image_shape 30 | self.k = k 31 | self.sigma = sigma 32 | self.method = method 33 | 34 | self.kernel = (2 * round(4 * sigma) + 1,) * 2 35 | 36 | def call(self, inputs): 37 | b = tf.shape(inputs)[0] 38 | h, w, c = tf.keras.backend.int_shape(inputs)[1:] 39 | feature = tf.reshape(inputs, [b * h * w, c]) 40 | score = tf.sort(euclidean_matrix(feature, self.feature_vector), axis = -1)[..., :self.k] 41 | mask = tf.reshape(score[..., 0], [b, h, w, 1]) 42 | score = tf.reshape(score, [b, h * w, -1]) 43 | #conf = tf.gather_nd(score, tf.stack([tf.range(b), tf.cast(tf.argmax(score[..., 0], axis = -1), tf.int32)], axis = -1)) 44 | conf = tf.gather(score, tf.argmax(score[..., 0], axis = -1), batch_dims = 1) 45 | exp_conf = tf.exp(conf) 46 | weight = 1 - tf.reduce_max(exp_conf, axis = -1) / tf.reduce_sum(exp_conf, axis = -1) 47 | score = tf.reduce_max(score[..., 0], axis = -1) * weight 48 | score = tf.expand_dims(score, axis = -1) 49 | 50 | #upsampling 51 | mask = tf.image.resize(mask, self.image_shape, method = self.method) 52 | 53 | #gaussian smoothing 54 | if 0 < self.sigma: 55 | mask = tf.py_function(lambda *args: np.array([cv2.GaussianBlur(m, self.kernel, self.sigma) for m in args[0].numpy()]), inp = [mask], Tout = mask.dtype) 56 | mask = tf.reshape(mask, [-1, *self.image_shape, 1]) 57 | return score, mask 58 | 59 | def get_config(self): 60 | config = super(Head, self).get_config() 61 | config["k"] = self.k 62 | config["image_shape"] = self.image_shape 63 | config["sigma"] = self.sigma 64 | config["method"] = self.method 65 | return config 66 | 67 | def patch_core_head(feature, feature_vector = None, image_shape = [224, 224], k = 9, sampling_index = None, pool_size = 3, sigma = 4, method = "bilinear", memory_reduce = False): 68 | feature = FeatureExtractor(sampling_index = sampling_index, pool_size = pool_size, memory_reduce = memory_reduce, name = "feature_extractor")(feature) 69 | if feature_vector is not None: 70 | score, mask = Head(feature_vector = feature_vector, image_shape = image_shape, k = k, sigma = sigma, method = method, name = "patch_core")(feature) 71 | return score, mask 72 | else: 73 | return feature -------------------------------------------------------------------------------- /tfdet/model/head/pspnet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.ops import AdaptiveAveragePooling2D, AdaptiveMaxPooling2D 4 | 5 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 6 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 7 | 8 | def head_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def normalize(axis = -1, momentum = 0.9, epsilon = 1e-5, **kwargs): 12 | return tf.keras.layers.BatchNormalization(axis = axis, momentum = momentum, epsilon = epsilon, **kwargs) 13 | 14 | class PoolingPyramidModule(tf.keras.layers.Layer): 15 | def __init__(self, pool_scale = [1, 2, 3, 6], n_feature = 512, max_pooling = False, method = "bilinear", 16 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu, **kwargs): 17 | super(PoolingPyramidModule, self).__init__(**kwargs) 18 | self.pool_scale = [pool_scale] if not isinstance(pool_scale, (tuple, list)) else pool_scale 19 | self.n_feature = n_feature 20 | self.max_pooling = max_pooling 21 | self.method = method 22 | self.convolution = convolution 23 | self.normalize = normalize 24 | self.activation = activation 25 | 26 | self.layers = [] 27 | for scale in self.pool_scale: 28 | layer = [(AdaptiveMaxPooling2D if max_pooling else AdaptiveAveragePooling2D)(scale)] 29 | layer.append(self.convolution(self.n_feature, 1, use_bias = self.normalize is None)) 30 | if self.normalize is not None: 31 | if scale == 1: #for bn moving_varriance nan bug fix 32 | layer.append(tf.keras.layers.Reshape([self.n_feature])) 33 | layer.append(self.normalize()) 34 | layer.append(tf.keras.layers.Reshape([1, 1, self.n_feature])) 35 | else: 36 | layer.append(self.normalize()) 37 | layer.append(tf.keras.layers.Activation(self.activation)) 38 | self.layers.append(layer) 39 | self.upsample = tf.keras.layers.Lambda(lambda args: tf.image.resize(args[0], args[1], method = method), name = "upsample") 40 | 41 | def call(self, inputs): 42 | target_size = tf.shape(inputs)[1:3] 43 | out = [] 44 | for layer in self.layers: 45 | o = inputs 46 | for l in layer: 47 | o = l(o) 48 | o = self.upsample([o, target_size]) 49 | out.append(o) 50 | return out 51 | 52 | def get_config(self): 53 | config = super(PoolingPyramidModule, self).get_config() 54 | config["pool_scale"] = self.pool_scale 55 | config["n_feature"] = self.n_feature 56 | config["max_pooling"] = self.max_pooling 57 | config["method"] = self.method 58 | return config 59 | 60 | def pspnet_head(feature, n_class = 35, n_feature = 512, pool_scale = [1, 2, 3, 6], max_pooling = False, method = "bilinear", dropout_rate = 0.1, 61 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 62 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 63 | #https://arxiv.org/abs/1612.01105 64 | if isinstance(feature, (tuple, list)): 65 | feature = feature[-1] 66 | out = PoolingPyramidModule(pool_scale, n_feature, max_pooling = max_pooling, method = method, convolution = convolution, normalize = normalize, activation = activation, name = "pooling_pyramoid_feature")(feature) 67 | out = tf.keras.layers.Concatenate(axis = -1, name = "feature_concat")([feature] + out) 68 | 69 | out = convolution(n_feature, 3, padding = "same", use_bias = normalize is None, name = "feature_conv")(out) 70 | if normalize is not None: 71 | out = normalize(name = "feature_norm")(out) 72 | out = tf.keras.layers.Activation(activation, name = "feature_act")(out) 73 | 74 | if 0 < dropout_rate: 75 | out = tf.keras.layers.Dropout(dropout_rate)(out) 76 | out = logits_convolution(n_class, 1, use_bias = True, name = "logits")(out) 77 | out = tf.keras.layers.Activation(logits_activation if logits_activation is not None else tf.keras.activations.linear, dtype = tf.float32, name = "logits_act")(out) 78 | return out -------------------------------------------------------------------------------- /tfdet/model/head/spade.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import cv2 3 | import numpy as np 4 | 5 | from tfdet.core.ops import feature_extract, euclidean, euclidean_matrix 6 | 7 | class FeatureExtractor(tf.keras.layers.Layer): 8 | def __init__(self, sampling_index = None, **kwargs): 9 | super(FeatureExtractor, self).__init__(**kwargs) 10 | self.sampling_index = sampling_index 11 | 12 | def call(self, inputs): 13 | out = feature_extract(inputs, sampling_index = self.sampling_index, pool_size = 1, sub_sampling = True, concat = False) 14 | return out 15 | 16 | def get_config(self): 17 | config = super(FeatureExtractor, self).get_config() 18 | config["sampling_index"] = self.sampling_index 19 | return config 20 | 21 | class Head(tf.keras.layers.Layer): 22 | def __init__(self, feature_vector, image_shape = [224, 224], k = 50, sigma = 4, method = "bilinear", **kwargs): 23 | super(Head, self).__init__(**kwargs) 24 | self.feature_vector = feature_vector 25 | self.image_shape = image_shape 26 | self.k = k 27 | self.sigma = sigma 28 | self.method = method 29 | 30 | self.kernel = (2 * round(4 * sigma) + 1,) * 2 31 | 32 | def call(self, inputs): 33 | #calculate image score with extract top-k feature by euclidean distance for gap feature 34 | dist = euclidean_matrix(tf.squeeze(inputs[-1], axis = [1, 2]), tf.squeeze(self.feature_vector[-1], axis = [1, 2])) 35 | score = tf.sort(dist, axis = -1)[..., :self.k] 36 | score = tf.reduce_mean(score, axis = -1, keepdims = True) 37 | indices = tf.argsort(dist, axis = -1)[..., :self.k] 38 | 39 | #calculate pixel score with extract top-k feature by euclidean distance for feature 40 | mask = [] 41 | for tr_feature, pred_feature in zip(self.feature_vector[:-1], inputs[:-1]): 42 | tr_feature = tf.gather(tr_feature, indices) 43 | pred_feature = tf.expand_dims(pred_feature, axis = 1) 44 | dist = euclidean(tr_feature, pred_feature) 45 | m = tf.reduce_min(dist, axis = 1) 46 | m = tf.image.resize(tf.expand_dims(m, axis = -1), self.image_shape, method = self.method) #upsampling 47 | mask.append(m) 48 | mask = tf.reduce_mean(mask, axis = 0) 49 | 50 | #gaussian smoothing 51 | if 0 < self.sigma: 52 | mask = tf.py_function(lambda *args: np.array([cv2.GaussianBlur(m, self.kernel, self.sigma) for m in args[0].numpy()]), inp = [mask], Tout = mask.dtype) 53 | mask = tf.reshape(mask, [-1, *self.image_shape, 1]) 54 | return score, mask 55 | 56 | def get_config(self): 57 | config = super(Head, self).get_config() 58 | config["k"] = self.k 59 | config["image_shape"] = self.image_shape 60 | config["sigma"] = self.sigma 61 | config["method"] = self.method 62 | return config 63 | 64 | def spade_head(feature, feature_vector = None, image_shape = [224, 224], k = 50, sampling_index = None, sigma = 4, method = "bilinear"): 65 | feature = FeatureExtractor(sampling_index = sampling_index, name = "feature_extractor")(feature) 66 | if feature_vector is not None: 67 | score, mask = Head(feature_vector = feature_vector, image_shape = image_shape, k = k, sigma = sigma, method = method, name = "spade")(feature) 68 | return score, mask 69 | else: 70 | return feature -------------------------------------------------------------------------------- /tfdet/model/head/upernet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from .pspnet import PoolingPyramidModule 4 | 5 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "he_normal", **kwargs): 6 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 7 | 8 | def head_conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = tf.keras.initializers.RandomNormal(mean = 0.0, stddev = 0.01), **kwargs): 9 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 10 | 11 | def normalize(axis = -1, momentum = 0.9, epsilon = 1e-5, **kwargs): 12 | return tf.keras.layers.BatchNormalization(axis = axis, momentum = momentum, epsilon = epsilon, **kwargs) 13 | 14 | def upernet_head(feature, n_class = 35, n_feature = 512, pool_scale = [1, 2, 3, 6], max_pooling = False, method = "bilinear", dropout_rate = 0.1, 15 | logits_convolution = head_conv, logits_activation = tf.keras.activations.softmax, 16 | convolution = conv, normalize = normalize, activation = tf.keras.activations.relu): 17 | #https://arxiv.org/abs/1807.10221 18 | feature = [feature] if not isinstance(feature, (tuple, list)) else list(feature) 19 | 20 | out = [] 21 | for i, x in enumerate(feature[:-1]): 22 | o = convolution(n_feature, 1, use_bias = normalize is None, name = "feature{0}_resample_conv".format(i + 1) if 1 < len(feature) else "feature_resample_conv")(x) 23 | if normalize is not None: 24 | o = normalize(name = "feature{0}_resample_norm".format(i + 1) if 1 < len(feature) else "feature_resample_norm")(o) 25 | o = tf.keras.layers.Activation(activation, name = "feature{0}_resample_act".format(i + 1) if 1 < len(feature) else "feature_resample_act")(o) 26 | out.append(o) 27 | 28 | psp_feature = PoolingPyramidModule(pool_scale, n_feature, max_pooling = max_pooling, method = method, convolution = convolution, normalize = normalize, activation = activation, name = "psp_feature_extract")(feature[-1]) 29 | psp_feature = tf.keras.layers.Concatenate(axis = -1, name = "psp_feature")([feature[-1]] + psp_feature) 30 | psp_feature = convolution(n_feature, 3, padding = "same", use_bias = normalize is None, name = "psp_feature_conv")(psp_feature) 31 | if normalize is not None: 32 | psp_feature = normalize(name = "psp_feature_norm")(psp_feature) 33 | psp_feature = tf.keras.layers.Activation(activation, name = "psp_feature_act")(psp_feature) 34 | out.append(psp_feature) 35 | 36 | upsample = tf.keras.layers.Lambda(lambda args: tf.image.resize(args[0], args[1], method = method), name = "feature_upsample") 37 | for index in range(len(out) - 1, 0, -1): 38 | prev_size = tf.shape(out[index - 1])[1:3] 39 | out[index - 1] = tf.keras.layers.Add(name = "feature{0}_add".format(index))([out[index - 1], upsample([out[index], prev_size])]) 40 | 41 | for index in range(len(out) - 1): #without psp feature 42 | o = convolution(n_feature, 3, padding = "same", use_bias = normalize is None, name = "feature{0}_post_conv".format(index + 1))(out[index]) 43 | if normalize is not None: 44 | o = normalize(name = "feature{0}_post_norm".format(index + 1))(o) 45 | o = tf.keras.layers.Activation(activation, name = "feature{0}_post_act".format(index + 1))(o) 46 | out[index] = o 47 | 48 | target_size = tf.shape(out[0])[1:3] 49 | upsample = tf.keras.layers.Lambda(lambda args: tf.image.resize(args[0], args[1], method = method), name = "feature_post_upsample") 50 | out = [upsample([o, target_size]) if i != 0 else o for i, o in enumerate(out)] 51 | out = tf.keras.layers.Concatenate(axis = -1, name = "feature_post_concat")(out) 52 | 53 | out = convolution(n_feature, 3, padding = "same", use_bias = normalize is None, name = "head_conv")(out) 54 | if normalize is not None: 55 | out = normalize(name = "head_norm")(out) 56 | out = tf.keras.layers.Activation(activation, name = "head_act")(out) 57 | 58 | if 0 < dropout_rate: 59 | out = tf.keras.layers.Dropout(dropout_rate)(out) 60 | out = logits_convolution(n_class, 1, use_bias = True, name = "logits")(out) 61 | out = tf.keras.layers.Activation(logits_activation if logits_activation is not None else tf.keras.activations.linear, dtype = tf.float32, name = "logits_act")(out) 62 | return out -------------------------------------------------------------------------------- /tfdet/model/neck/__init__.py: -------------------------------------------------------------------------------- 1 | from .align import FeatureAlign, FeatureUpsample 2 | from .fpn import fpn, panet, bifpn, FeaturePyramidNetwork -------------------------------------------------------------------------------- /tfdet/model/neck/align.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from .fpn import fpn 4 | 5 | def conv(filters, kernel_size, strides = 1, padding = "same", use_bias = True, kernel_initializer = "glorot_uniform", **kwargs): 6 | return tf.keras.layers.Conv2D(filters, kernel_size, strides = strides, padding = padding, use_bias = use_bias, kernel_initializer = kernel_initializer, **kwargs) 7 | 8 | class FeatureAlign(tf.keras.layers.Layer): 9 | def __init__(self, n_feature = 256, n_sampling = 1, pre_sampling = True, neck = fpn, neck_n_depth = 1, use_bias = None, convolution = conv, normalize = None, **kwargs): 10 | super(FeatureAlign, self).__init__(**kwargs) 11 | self.n_feature = n_feature 12 | self.n_sampling = n_sampling 13 | self.pre_sampling = pre_sampling 14 | self.neck = neck 15 | self.neck_n_depth = neck_n_depth 16 | self.use_bias = (normalize is None) if use_bias is None else use_bias 17 | self.convolution = convolution 18 | self.normalize = normalize 19 | 20 | def build(self, input_shape): 21 | if not isinstance(input_shape, (tuple, list)): 22 | input_shape = [input_shape] 23 | 24 | self.layers = [] 25 | for index in range(self.n_sampling): 26 | layer = [] 27 | if index == 0: 28 | layer.append(self.convolution(self.n_feature, 1, use_bias = self.use_bias)) 29 | if self.normalize is not None: 30 | layer.append(self.normalize()) 31 | layer.append(tf.keras.layers.MaxPool2D(3, strides = 2, padding = "same")) 32 | self.layers.append(layer) 33 | 34 | self.neck_layers = [] 35 | if self.neck is not None and 0 < self.neck_n_depth: 36 | for index in range(self.neck_n_depth): 37 | self.neck_layers.append(self.neck()) 38 | else: 39 | self.neck_layers.append([self.convolution(self.n_feature, 1, use_bias = True) for index in range(len(input_shape) + (self.n_sampling if self.pre_sampling else 0))]) 40 | 41 | def call(self, inputs, **kwargs): 42 | if not isinstance(inputs, (tuple, list)): 43 | inputs = [inputs] 44 | feature = list(inputs) 45 | 46 | if self.pre_sampling: 47 | x = feature[-1] 48 | for layers in self.layers: 49 | for l in layers: 50 | x = l(x) 51 | feature.append(x) 52 | 53 | for layer in self.neck_layers: 54 | if not isinstance(layer, (tuple, list)): 55 | feature = layer(feature) 56 | else: 57 | print(layer) 58 | for index in range(len(feature)): 59 | feature[index] = layer[index](feature[index]) 60 | 61 | if not isinstance(feature, (tuple, list)): 62 | feature = [feature] 63 | if not self.pre_sampling: 64 | x = feature[-1] 65 | for layers in self.layers: 66 | for l in layers: 67 | x = l(x) 68 | feature.append(x) 69 | 70 | if len(feature) == 1: 71 | feature = feature[0] 72 | return feature 73 | 74 | def get_config(self): 75 | config = super(FeatureAlign, self).get_config() 76 | config["n_feature"] = self.n_feature 77 | config["n_sampling"] = self.n_sampling 78 | config["pre_sampling"] = self.pre_sampling 79 | config["neck_n_depth"] = self.neck_n_depth 80 | config["use_bias"] = self.use_bias 81 | return config 82 | 83 | class FeatureUpsample(tf.keras.layers.Layer): 84 | def __init__(self, concat = True, method = "bilinear", **kwargs): 85 | super(FeatureUpsample, self).__init__(**kwargs) 86 | self.concat = concat 87 | self.method = method 88 | self.upsample = tf.keras.layers.Lambda(lambda args: tf.image.resize(args[0], args[1], method = method), name = "upsample") 89 | if self.concat: 90 | self.post = tf.keras.layers.Concatenate(axis = -1) 91 | 92 | def call(self, inputs): 93 | if not isinstance(inputs, (tuple, list)): 94 | inputs = [inputs] 95 | 96 | target_size = tf.shape(inputs[0])[1:3] 97 | out = [inputs[0]] + [self.upsample([x, target_size]) for x in inputs[1:]] 98 | 99 | if self.concat and 1 < len(out): 100 | out = self.post(out) 101 | elif len(out) == 1: 102 | out = out[0] 103 | return out 104 | 105 | def get_config(self): 106 | config = super(FeatureUpsample, self).get_config() 107 | config["concat"] = self.concat 108 | config["method"] = self.method 109 | return config -------------------------------------------------------------------------------- /tfdet/model/postprocess/__init__.py: -------------------------------------------------------------------------------- 1 | from . import anchor 2 | from . import anchor_free 3 | from . import roi 4 | from . import yolo 5 | effdet = retina = anchor 6 | fcos = anchor_free 7 | rcnn = roi 8 | 9 | from . import anodet 10 | spade = padim = patch_core = anodet -------------------------------------------------------------------------------- /tfdet/model/postprocess/anchor.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.bbox import delta2bbox 4 | from tfdet.core.ops import multiclass_nms 5 | from tfdet.core.util import map_fn 6 | 7 | class FilterDetection(tf.keras.layers.Layer): 8 | def __init__(self, proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, valid_inside_anchor = False, ignore_label = 0, performance_count = 5000, 9 | mean = [0., 0., 0., 0.], std = [1., 1., 1., 1.], clip_ratio = 16 / 1000, 10 | batch_size = 1, dtype = tf.float32, 11 | tensorrt = False, **kwargs): 12 | kwargs["dtype"] = dtype 13 | super(FilterDetection, self).__init__(**kwargs) 14 | self.proposal_count = proposal_count 15 | self.iou_threshold = iou_threshold 16 | self.score_threshold = score_threshold 17 | self.soft_nms = soft_nms 18 | self.valid_inside_anchor = valid_inside_anchor 19 | self.ignore_label = ignore_label 20 | self.performance_count = performance_count 21 | self.mean = mean 22 | self.std = std 23 | self.clip_ratio = clip_ratio 24 | self.batch_size = batch_size 25 | self.tensorrt = tensorrt 26 | 27 | def call(self, inputs): 28 | y_pred, bbox_pred, anchors = inputs[:3] 29 | conf_pred = inputs[3] if 3 < len(inputs) else None 30 | if isinstance(y_pred, (list, tuple)): 31 | y_pred = tf.concat(y_pred, axis = -2) 32 | if isinstance(bbox_pred, (list, tuple)): 33 | bbox_pred = tf.concat(bbox_pred, axis = -2) 34 | if isinstance(anchors, (list, tuple)): 35 | anchors = tf.concat(anchors, axis = 0) 36 | if conf_pred is not None and isinstance(conf_pred, (list, tuple)): 37 | conf_pred = tf.concat(conf_pred, axis = -2) 38 | 39 | if self.valid_inside_anchor: 40 | valid_flags = tf.logical_and(tf.less_equal(anchors[..., 2], 1), 41 | tf.logical_and(tf.less_equal(anchors[..., 3], 1), 42 | tf.logical_and(tf.greater_equal(anchors[..., 0], 0), 43 | tf.greater_equal(anchors[..., 1], 0)))) 44 | #valid_indices = tf.range(tf.shape(anchors)[0])[valid_flags] 45 | valid_indices = tf.where(valid_flags)[:, 0] 46 | y_pred = tf.gather(y_pred, valid_indices, axis = 1) 47 | bbox_pred = tf.gather(bbox_pred, valid_indices, axis = 1) 48 | anchors = tf.gather(anchors, valid_indices) 49 | if conf_pred is not None: 50 | conf_pred = tf.gather(conf_pred, valid_indices, axis = 1) 51 | if conf_pred is not None: 52 | y_pred = tf.multiply(y_pred, conf_pred) 53 | 54 | if not self.tensorrt: 55 | anchors = tf.tile(tf.expand_dims(anchors, axis = 0), [tf.shape(y_pred)[0], 1, 1]) 56 | out = map_fn(multiclass_nms, y_pred, bbox_pred, anchors, dtype = (self.dtype, self.dtype), batch_size = self.batch_size, 57 | proposal_count = self.proposal_count, soft_nms = self.soft_nms, iou_threshold = self.iou_threshold, score_threshold = self.score_threshold, ignore_label = self.ignore_label, performance_count = self.performance_count, 58 | coder_func = delta2bbox, mean = self.mean, std = self.std, clip_ratio = self.clip_ratio) 59 | else: 60 | bbox_pred = delta2bbox(anchors, bbox_pred, mean = self.mean, std = self.std, clip_ratio = self.clip_ratio) 61 | bbox_pred = tf.clip_by_value(bbox_pred, 0, 1) 62 | out = (y_pred, bbox_pred) 63 | return out 64 | 65 | def get_config(self): 66 | config = super(FilterDetection, self).get_config() 67 | config["proposal_count"] = self.proposal_count 68 | config["iou_threshold"] = self.iou_threshold 69 | config["score_threshold"] = self.score_threshold 70 | config["soft_nms"] = self.soft_nms 71 | config["valid_inside_anchor"] = self.valid_inside_anchor 72 | config["ignore_label"] = self.ignore_label 73 | config["performance_count"] = self.performance_count 74 | config["mean"] = self.mean 75 | config["std"] = self.std 76 | config["clip_ratio"] = self.clip_ratio 77 | config["batch_size"] = self.batch_size 78 | config["tensorrt"] = self.tensorrt 79 | return config -------------------------------------------------------------------------------- /tfdet/model/postprocess/anchor_free.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.bbox import offset2bbox 4 | from tfdet.core.ops import multiclass_nms 5 | from tfdet.core.util import map_fn 6 | 7 | class FilterDetection(tf.keras.layers.Layer): 8 | def __init__(self, proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, ignore_label = 0, performance_count = 5000, 9 | batch_size = 1, dtype = tf.float32, 10 | tensorrt = False, **kwargs): 11 | kwargs["dtype"] = dtype 12 | super(FilterDetection, self).__init__(**kwargs) 13 | self.proposal_count = proposal_count 14 | self.iou_threshold = iou_threshold 15 | self.score_threshold = score_threshold 16 | self.soft_nms = soft_nms 17 | self.ignore_label = ignore_label 18 | self.performance_count = performance_count 19 | self.batch_size = batch_size 20 | self.tensorrt = tensorrt 21 | 22 | def call(self, inputs): 23 | y_pred, bbox_pred, points = inputs[:3] 24 | conf_pred = inputs[3] if 3 < len(inputs) else None 25 | if isinstance(y_pred, (list, tuple)): 26 | y_pred = tf.concat(y_pred, axis = -2) 27 | if isinstance(bbox_pred, (list, tuple)): 28 | bbox_pred = tf.concat(bbox_pred, axis = -2) 29 | if isinstance(points, (list, tuple)): 30 | points = tf.concat(points, axis = 0) 31 | if conf_pred is not None: 32 | if isinstance(conf_pred, (list, tuple)): 33 | conf_pred = tf.concat(conf_pred, axis = -2) 34 | y_pred = tf.multiply(y_pred, conf_pred) 35 | y_pred = tf.sqrt(y_pred) 36 | 37 | if not self.tensorrt: 38 | points = tf.tile(tf.expand_dims(points, axis = 0), [tf.shape(y_pred)[0], 1, 1]) 39 | out = map_fn(multiclass_nms, y_pred, bbox_pred, points, dtype = (self.dtype, self.dtype), batch_size = self.batch_size, 40 | proposal_count = self.proposal_count, soft_nms = self.soft_nms, iou_threshold = self.iou_threshold, score_threshold = self.score_threshold, ignore_label = self.ignore_label, performance_count = self.performance_count, coder_func = offset2bbox) 41 | else: 42 | bbox_pred = offset2bbox(points, bbox_pred) 43 | bbox_pred = tf.clip_by_value(bbox_pred, 0, 1) 44 | out = (y_pred, bbox_pred) 45 | return out 46 | 47 | def get_config(self): 48 | config = super(FilterDetection, self).get_config() 49 | config["proposal_count"] = self.proposal_count 50 | config["iou_threshold"] = self.iou_threshold 51 | config["score_threshold"] = self.score_threshold 52 | config["soft_nms"] = self.soft_nms 53 | config["ignore_label"] = self.ignore_label 54 | config["performance_count"] = self.performance_count 55 | config["batch_size"] = self.batch_size 56 | config["tensorrt"] = self.tensorrt 57 | return config -------------------------------------------------------------------------------- /tfdet/model/postprocess/anodet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import cv2 3 | import numpy as np 4 | 5 | class FilterDetection(tf.keras.layers.Layer): 6 | def __init__(self, threshold, kernel_size = 4, **kwargs): 7 | super(FilterDetection, self).__init__(**kwargs) 8 | self.threshold = threshold 9 | self.kernel_size = kernel_size 10 | 11 | def call(self, inputs): 12 | score, mask = inputs 13 | score = tf.where(self.threshold <= score, score, 0) 14 | mask = tf.where(self.threshold <= mask, mask, 0) 15 | mask = tf.py_function(lambda *args: np.array([cv2.morphologyEx(o, cv2.MORPH_OPEN, args[1].numpy()) for o in args[0].numpy()]), inp = [mask, tf.ones([self.kernel_size] * 2, dtype = tf.int32)], Tout = mask.dtype) 16 | mask = tf.reshape(mask, tf.shape(inputs[1])) 17 | return score, mask 18 | 19 | def get_config(self): 20 | config = super(FilterDetection, self).get_config() 21 | config["threshold"] = self.threshold 22 | config["kernel_size"] = self.kernel_size 23 | return config -------------------------------------------------------------------------------- /tfdet/model/postprocess/fcos.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.bbox import offset2bbox 4 | from tfdet.core.util import map_fn, multiclass_nms 5 | 6 | class FilterDetection(tf.keras.layers.Layer): 7 | def __init__(self, proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, ignore_label = 0, performance_count = 5000, 8 | batch_size = 1, dtype = tf.float32, 9 | tensorrt = False, **kwargs): 10 | kwargs["dtype"] = dtype 11 | super(FilterDetection, self).__init__(**kwargs) 12 | self.proposal_count = proposal_count 13 | self.iou_threshold = iou_threshold 14 | self.score_threshold = score_threshold 15 | self.soft_nms = soft_nms 16 | self.ignore_label = ignore_label 17 | self.performance_count = performance_count 18 | self.batch_size = batch_size 19 | self.tensorrt = tensorrt 20 | 21 | def call(self, inputs): 22 | logits, regress, points = inputs[:3] 23 | centerness = inputs[3] if 3 < len(inputs) else None 24 | 25 | if isinstance(logits, list): 26 | logits = tf.concat(logits, axis = -2) 27 | regress = tf.concat(regress, axis = -2) 28 | points = tf.concat(points, axis = 0) 29 | if centerness is not None: 30 | centerness = tf.concat(centerness, axis = -2) 31 | if centerness is not None: 32 | logits = tf.multiply(logits, centerness) 33 | logits = tf.sqrt(logits) 34 | 35 | if not self.tensorrt: 36 | points = tf.tile(tf.expand_dims(points, axis = 0), [tf.shape(logits)[0], 1, 1]) 37 | out = map_fn(multiclass_nms, logits, regress, points, dtype = (self.dtype, self.dtype), batch_size = self.batch_size, 38 | proposal_count = self.proposal_count, soft_nms = self.soft_nms, iou_threshold = self.iou_threshold, score_threshold = self.score_threshold, ignore_label = self.ignore_label, performance_count = self.performance_count, coder_func = offset2bbox) 39 | else: 40 | regress = offset2bbox(points, regress) 41 | regress = tf.clip_by_value(regress, 0, 1) 42 | out = (logits, regress) 43 | return out 44 | 45 | def get_config(self): 46 | config = super(FilterDetection, self).get_config() 47 | config["proposal_count"] = self.proposal_count 48 | config["iou_threshold"] = self.iou_threshold 49 | config["score_threshold"] = self.score_threshold 50 | config["soft_nms"] = self.soft_nms 51 | config["ignore_label"] = self.ignore_label 52 | config["performance_count"] = self.performance_count 53 | config["batch_size"] = self.batch_size 54 | config["tensorrt"] = self.tensorrt 55 | return config -------------------------------------------------------------------------------- /tfdet/model/postprocess/rcnn.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from tfdet.core.bbox import delta2bbox 5 | from tfdet.core.util import map_fn, multiclass_nms 6 | 7 | class FilterDetection(tf.keras.layers.Layer): 8 | def __init__(self, proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, ensemble = True, valid = False, ignore_label = 0, performance_count = 5000, 9 | batch_size = 1, mean = [0., 0., 0., 0.], std = [0.1, 0.1, 0.2, 0.2], clip_ratio = 16 / 1000, dtype = tf.float32, 10 | tensorrt = False, **kwargs): 11 | kwargs["dtype"] = dtype 12 | super(FilterDetection, self).__init__(**kwargs) 13 | self.proposal_count = proposal_count 14 | self.iou_threshold = iou_threshold 15 | self.score_threshold = score_threshold 16 | self.soft_nms = soft_nms 17 | self.ensemble = ensemble 18 | self.valid = valid 19 | self.ignore_label = ignore_label 20 | self.performance_count = performance_count 21 | self.batch_size = batch_size 22 | self.mean = mean 23 | self.std = std 24 | self.clip_ratio = clip_ratio 25 | self.tensorrt = tensorrt 26 | 27 | def call(self, inputs): 28 | if 5 < len(inputs): 29 | inputs = inputs[3:7] 30 | cls_logits, cls_regress, proposals = inputs[:3] 31 | mask_regress = inputs[3] if 3 < len(inputs) else None 32 | 33 | std = self.std 34 | if isinstance(proposals, list): 35 | proposals = proposals[-1] 36 | if isinstance(cls_logits, list): 37 | std = np.divide(std, len(cls_logits)) 38 | cls_logits, cls_regress = tf.reduce_mean(cls_logits, axis = 0) if self.ensemble else cls_logits[-1], cls_regress[-1] 39 | if mask_regress is not None: 40 | mask_regress = tf.reduce_mean(mask_regress, axis = 0) if self.ensemble else mask_regress[-1] 41 | elif tf.keras.backend.int_shape(cls_logits)[-1] == 1: #rpn_score, rpn_regress, anchors 42 | cls_logits = tf.concat([1 - cls_logits, cls_logits], axis = -1) 43 | if self.valid: 44 | valid_flags = tf.logical_and(tf.less_equal(proposals[..., 2], 1), 45 | tf.logical_and(tf.less_equal(proposals[..., 3], 1), 46 | tf.logical_and(tf.greater_equal(proposals[..., 0], 0), 47 | tf.greater_equal(proposals[..., 1], 0)))) 48 | #valid_indices = tf.range(tf.shape(proposals)[0])[valid_flags] 49 | valid_indices = tf.where(valid_flags)[:, 0] 50 | cls_logits = tf.gather(cls_logits, valid_indices, axis = 1) 51 | cls_regress = tf.gather(cls_regress, valid_indices, axis = 1) 52 | proposals = tf.gather(proposals, valid_indices) 53 | if not self.tensorrt: 54 | proposals = tf.tile(tf.expand_dims(proposals, axis = 0), [tf.shape(cls_logits)[0], 1, 1]) 55 | 56 | dtype = (self.dtype, self.dtype) 57 | if mask_regress is not None: 58 | dtype += (self.dtype,) 59 | args = [l for l in [cls_logits, cls_regress, proposals, mask_regress] if l is not None] 60 | if not self.tensorrt: 61 | out = map_fn(multiclass_nms, *args, dtype = dtype, batch_size = self.batch_size, 62 | proposal_count = self.proposal_count, iou_threshold = self.iou_threshold, score_threshold = self.score_threshold, soft_nms = self.soft_nms, ignore_label = self.ignore_label, performance_count = self.performance_count, 63 | coder_func = delta2bbox, mean = self.mean, std = std, clip_ratio = self.clip_ratio) 64 | else: 65 | raise ValueError("Conversion of rcnn is not yet supported.") 66 | return out 67 | 68 | def get_config(self): 69 | config = super(FilterDetection, self).get_config() 70 | config["proposal_count"] = self.proposal_count 71 | config["iou_threshold"] = self.iou_threshold 72 | config["score_threshold"] = self.score_threshold 73 | config["soft_nms"] = self.soft_nms 74 | config["ensemble"] = self.ensemble 75 | config["valid"] = self.valid 76 | config["ignore_label"] = self.ignore_label 77 | config["performance_count"] = self.performance_count 78 | config["batch_size"] = self.batch_size 79 | config["mean"] = self.mean 80 | config["std"] = self.std 81 | config["clip_ratio"] = self.clip_ratio 82 | config["tensorrt"] = self.tensorrt 83 | return config -------------------------------------------------------------------------------- /tfdet/model/postprocess/retina.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.bbox import delta2bbox 4 | from tfdet.core.util import map_fn, multiclass_nms 5 | 6 | class FilterDetection(tf.keras.layers.Layer): 7 | def __init__(self, proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, valid = False, ignore_label = 0, performance_count = 5000, 8 | batch_size = 1, mean = [0., 0., 0., 0.], std = [1., 1., 1., 1.], clip_ratio = 16 / 1000, dtype = tf.float32, 9 | tensorrt = False, **kwargs): 10 | kwargs["dtype"] = dtype 11 | super(FilterDetection, self).__init__(**kwargs) 12 | self.proposal_count = proposal_count 13 | self.iou_threshold = iou_threshold 14 | self.score_threshold = score_threshold 15 | self.soft_nms = soft_nms 16 | self.valid = valid 17 | self.ignore_label = ignore_label 18 | self.performance_count = performance_count 19 | self.batch_size = batch_size 20 | self.mean = mean 21 | self.std = std 22 | self.clip_ratio = clip_ratio 23 | self.tensorrt = tensorrt 24 | 25 | def call(self, inputs): 26 | logits, regress, anchors = inputs[:3] 27 | confs = inputs[3] if 3 < len(inputs) else None 28 | 29 | if self.valid: 30 | valid_flags = tf.logical_and(tf.less_equal(anchors[..., 2], 1), 31 | tf.logical_and(tf.less_equal(anchors[..., 3], 1), 32 | tf.logical_and(tf.greater_equal(anchors[..., 0], 0), 33 | tf.greater_equal(anchors[..., 1], 0)))) 34 | #valid_indices = tf.range(tf.shape(anchors)[0])[valid_flags] 35 | valid_indices = tf.where(valid_flags)[:, 0] 36 | logits = tf.gather(logits, valid_indices, axis = 1) 37 | regress = tf.gather(regress, valid_indices, axis = 1) 38 | anchors = tf.gather(anchors, valid_indices) 39 | if confs is not None: 40 | confs = tf.gather(confs, valid_indices, axis = 1) 41 | if confs is not None: 42 | logits = tf.multiply(logits, confs) 43 | 44 | if not self.tensorrt: 45 | anchors = tf.tile(tf.expand_dims(anchors, axis = 0), [tf.shape(logits)[0], 1, 1]) 46 | out = map_fn(multiclass_nms, logits, regress, anchors, dtype = (self.dtype, self.dtype), batch_size = self.batch_size, 47 | proposal_count = self.proposal_count, soft_nms = self.soft_nms, iou_threshold = self.iou_threshold, score_threshold = self.score_threshold, ignore_label = self.ignore_label, performance_count = self.performance_count, 48 | coder_func = delta2bbox, mean = self.mean, std = self.std, clip_ratio = self.clip_ratio) 49 | else: 50 | regress = delta2bbox(anchors, regress, mean = self.mean, std = self.std, clip_ratio = self.clip_ratio) 51 | regress = tf.clip_by_value(regress, 0, 1) 52 | out = (logits, regress) 53 | return out 54 | 55 | def get_config(self): 56 | config = super(FilterDetection, self).get_config() 57 | config["proposal_count"] = self.proposal_count 58 | config["iou_threshold"] = self.iou_threshold 59 | config["score_threshold"] = self.score_threshold 60 | config["soft_nms"] = self.soft_nms 61 | config["valid"] = self.valid 62 | config["ignore_label"] = self.ignore_label 63 | config["performance_count"] = self.performance_count 64 | config["batch_size"] = self.batch_size 65 | config["mean"] = self.mean 66 | config["std"] = self.std 67 | config["clip_ratio"] = self.clip_ratio 68 | config["tensorrt"] = self.tensorrt 69 | return config -------------------------------------------------------------------------------- /tfdet/model/postprocess/roi.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from tfdet.core.bbox import delta2bbox 5 | from tfdet.core.ops import multiclass_nms 6 | from tfdet.core.util import map_fn 7 | 8 | class FilterDetection(tf.keras.layers.Layer): 9 | def __init__(self, proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, ensemble = True, valid_inside_anchor = False, ignore_label = 0, performance_count = 5000, 10 | mean = [0., 0., 0., 0.], std = [0.1, 0.1, 0.2, 0.2], clip_ratio = 16 / 1000, 11 | batch_size = 1, dtype = tf.float32, 12 | tensorrt = False, **kwargs): 13 | kwargs["dtype"] = dtype 14 | super(FilterDetection, self).__init__(**kwargs) 15 | self.proposal_count = proposal_count 16 | self.iou_threshold = iou_threshold 17 | self.score_threshold = score_threshold 18 | self.soft_nms = soft_nms 19 | self.ensemble = ensemble 20 | self.valid_inside_anchor = valid_inside_anchor 21 | self.ignore_label = ignore_label 22 | self.performance_count = performance_count 23 | self.mean = mean 24 | self.std = std 25 | self.clip_ratio = clip_ratio 26 | self.batch_size = batch_size 27 | self.tensorrt = tensorrt 28 | 29 | def call(self, inputs): 30 | if 5 < len(inputs): 31 | inputs = inputs[3:] 32 | y_pred, bbox_pred, proposals = inputs[:3] 33 | mask_pred = inputs[3] if 3 < len(inputs) else None 34 | if mask_pred is not None and tf.keras.backend.ndim(mask_pred[0] if isinstance(mask_pred, (tuple, list)) else mask_pred) != 5: 35 | mask_pred = None 36 | 37 | n_class = tf.keras.backend.int_shape(y_pred[0] if isinstance(y_pred, (tuple, list)) else y_pred)[-1] 38 | std = self.std 39 | if 1 < n_class: #2-stage cls 40 | if isinstance(proposals, (list, tuple)): 41 | proposals = proposals[-1] 42 | if isinstance(y_pred, (list, tuple)): 43 | std = np.divide(std, len(y_pred)) 44 | y_pred, bbox_pred = tf.reduce_mean(y_pred, axis = 0) if self.ensemble else y_pred[-1], bbox_pred[-1] 45 | if mask_pred is not None: 46 | mask_pred = tf.reduce_mean(mask_pred, axis = 0) if self.ensemble else mask_pred[-1] 47 | else: #1-stage rpn 48 | if isinstance(y_pred, (list, tuple)): 49 | y_pred = tf.concat(y_pred, axis = -2) 50 | if isinstance(bbox_pred, (list, tuple)): 51 | bbox_pred = tf.concat(bbox_pred, axis = -2) 52 | if isinstance(proposals, (list, tuple)): 53 | proposals = tf.concat(proposals, axis = 0) 54 | #y_pred = tf.concat([1 - y_pred, y_pred], axis = -1) 55 | if self.valid_inside_anchor: 56 | valid_flags = tf.logical_and(tf.less_equal(proposals[..., 2], 1), 57 | tf.logical_and(tf.less_equal(proposals[..., 3], 1), 58 | tf.logical_and(tf.greater_equal(proposals[..., 0], 0), 59 | tf.greater_equal(proposals[..., 1], 0)))) 60 | #valid_indices = tf.range(tf.shape(proposals)[0])[valid_flags] 61 | valid_indices = tf.where(valid_flags)[:, 0] 62 | y_pred = tf.gather(y_pred, valid_indices, axis = 1) 63 | bbox_pred = tf.gather(bbox_pred, valid_indices, axis = 1) 64 | proposals = tf.gather(proposals, valid_indices) 65 | if not self.tensorrt: 66 | proposals = tf.tile(tf.expand_dims(proposals, axis = 0), [tf.shape(y_pred)[0], 1, 1]) 67 | 68 | dtype = (self.dtype, self.dtype) 69 | if mask_pred is not None: 70 | dtype = (self.dtype, self.dtype, self.dtype) 71 | args = [l for l in [y_pred, bbox_pred, proposals, mask_pred] if l is not None] 72 | if not self.tensorrt: 73 | out = map_fn(multiclass_nms, *args, dtype = dtype, batch_size = self.batch_size, 74 | proposal_count = self.proposal_count, iou_threshold = self.iou_threshold, score_threshold = self.score_threshold, soft_nms = self.soft_nms, ignore_label = self.ignore_label, performance_count = self.performance_count, 75 | coder_func = delta2bbox, mean = self.mean, std = std, clip_ratio = self.clip_ratio) 76 | else: 77 | raise ValueError("Conversion of rcnn is not yet supported.") 78 | return out 79 | 80 | def get_config(self): 81 | config = super(FilterDetection, self).get_config() 82 | config["proposal_count"] = self.proposal_count 83 | config["iou_threshold"] = self.iou_threshold 84 | config["score_threshold"] = self.score_threshold 85 | config["soft_nms"] = self.soft_nms 86 | config["ensemble"] = self.ensemble 87 | config["valid_inside_anchor"] = self.valid_inside_anchor 88 | config["ignore_label"] = self.ignore_label 89 | config["performance_count"] = self.performance_count 90 | config["mean"] = self.mean 91 | config["std"] = self.std 92 | config["clip_ratio"] = self.clip_ratio 93 | config["batch_size"] = self.batch_size 94 | config["tensorrt"] = self.tensorrt 95 | return config -------------------------------------------------------------------------------- /tfdet/model/postprocess/yolo.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.bbox import yolo2bbox 4 | from tfdet.core.ops import multiclass_nms 5 | from tfdet.core.util import map_fn 6 | 7 | class FilterDetection(tf.keras.layers.Layer): 8 | def __init__(self, proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, valid_inside_anchor = False, ignore_label = 0, performance_count = 5000, 9 | clip_ratio = 16 / 1000, 10 | batch_size = 1, dtype = tf.float32, 11 | tensorrt = False, **kwargs): 12 | kwargs["dtype"] = dtype 13 | super(FilterDetection, self).__init__(**kwargs) 14 | self.proposal_count = proposal_count 15 | self.iou_threshold = iou_threshold 16 | self.score_threshold = score_threshold 17 | self.soft_nms = soft_nms 18 | self.valid_inside_anchor = valid_inside_anchor 19 | self.ignore_label = ignore_label 20 | self.performance_count = performance_count 21 | self.clip_ratio = clip_ratio 22 | self.batch_size = batch_size 23 | self.tensorrt = tensorrt 24 | 25 | def call(self, inputs): 26 | score_pred, logit_pred, bbox_pred, anchors = inputs 27 | if isinstance(score_pred, (list, tuple)): 28 | score_pred = tf.concat(score_pred, axis = -2) 29 | if isinstance(logit_pred, (list, tuple)): 30 | logit_pred = tf.concat(logit_pred, axis = -2) 31 | if isinstance(bbox_pred, (list, tuple)): 32 | bbox_pred = tf.concat(bbox_pred, axis = -2) 33 | if isinstance(anchors, (list, tuple)): 34 | anchors = tf.concat(anchors, axis = 0) 35 | 36 | if self.valid_inside_anchor: 37 | valid_flags = tf.logical_and(tf.less_equal(anchors[..., 2], 1), 38 | tf.logical_and(tf.less_equal(anchors[..., 3], 1), 39 | tf.logical_and(tf.greater_equal(anchors[..., 0], 0), 40 | tf.greater_equal(anchors[..., 1], 0)))) 41 | #valid_indices = tf.range(tf.shape(anchors)[0])[valid_flags] 42 | valid_indices = tf.where(valid_flags)[:, 0] 43 | score_pred = tf.gather(score_pred, valid_indices, axis = 1) 44 | logit_pred = tf.gather(logit_pred, valid_indices, axis = 1) 45 | bbox_pred = tf.gather(bbox_pred, valid_indices, axis = 1) 46 | anchors = tf.gather(anchors, valid_indices) 47 | y_pred = tf.multiply(logit_pred, score_pred) 48 | 49 | if not self.tensorrt: 50 | anchors = tf.tile(tf.expand_dims(anchors, axis = 0), [tf.shape(y_pred)[0], 1, 1]) 51 | out = map_fn(multiclass_nms, y_pred, bbox_pred, anchors, dtype = (self.dtype, self.dtype), batch_size = self.batch_size, 52 | proposal_count = self.proposal_count, iou_threshold = self.iou_threshold, score_threshold = self.score_threshold, soft_nms = self.soft_nms, ignore_label = self.ignore_label, performance_count = self.performance_count, 53 | coder_func = yolo2bbox, clip_ratio = self.clip_ratio) 54 | else: 55 | bbox_pred = yolo2bbox(anchors, bbox_pred, clip_ratio = self.clip_ratio) 56 | bbox_pred = tf.clip_by_value(bbox_pred, 0, 1) 57 | out = (y_pred, bbox_pred) 58 | return out 59 | 60 | def get_config(self): 61 | config = super(FilterDetection, self).get_config() 62 | config["proposal_count"] = self.proposal_count 63 | config["iou_threshold"] = self.iou_threshold 64 | config["score_threshold"] = self.score_threshold 65 | config["soft_nms"] = self.soft_nms 66 | config["valid_inside_anchor"] = self.valid_inside_anchor 67 | config["ignore_label"] = self.ignore_label 68 | config["performance_count"] = self.performance_count 69 | config["clip_ratio"] = self.clip_ratio 70 | config["batch_size"] = self.batch_size 71 | config["tensorrt"] = self.tensorrt 72 | return config 73 | -------------------------------------------------------------------------------- /tfdet/model/train/__init__.py: -------------------------------------------------------------------------------- 1 | from . import loss 2 | 3 | from . import fcos 4 | from . import rcnn 5 | from . import retina 6 | from . import yolo 7 | effdet = retina 8 | 9 | from . import spade 10 | from . import padim 11 | from . import patch_core 12 | 13 | from . import segmentation -------------------------------------------------------------------------------- /tfdet/model/train/fcos.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.assign import point 4 | from tfdet.core.loss import binary_cross_entropy, focal_binary_cross_entropy, iou, regularize as regularize_loss 5 | from .loss import AnchorFreeLoss 6 | from ..postprocess.anchor_free import FilterDetection 7 | 8 | def focal_loss(y_true, y_pred, alpha = .25, gamma = 2., weight = None, reduce = tf.reduce_mean): 9 | return focal_binary_cross_entropy(y_true, y_pred, alpha = alpha, gamma = gamma, weight = weight, reduce = reduce) 10 | 11 | def train_model(input, y_pred, bbox_pred, points, conf_pred = None, 12 | assign = point, sampler = None, 13 | proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, ignore_label = 0, performance_count = 5000, 14 | class_loss = focal_loss, bbox_loss = iou, conf_loss = binary_cross_entropy, 15 | regularize = True, weight_decay = 1e-4, 16 | decode_bbox = True, class_weight = None, background = False, 17 | batch_size = 1, missing_value = 0.): 18 | if not isinstance(y_pred, (tuple, list)): 19 | y_pred = [y_pred] 20 | if not isinstance(bbox_pred, (tuple, list)): 21 | bbox_pred = [bbox_pred] 22 | if not isinstance(points, (tuple, list)): 23 | points = [points] 24 | if conf_pred is not None and not isinstance(conf_pred, (tuple, list)): 25 | conf_pred = [conf_pred] 26 | 27 | y_true = tf.keras.layers.Input(shape = (None, None), name = "y_true") 28 | bbox_true = tf.keras.layers.Input(shape = (None, 4), name = "bbox_true") 29 | 30 | #[[0, 64], [64, 128], [128, 256], [256, 512], [512, float("inf")]] > [[0, 0.0625], [0.0625, 0.125], [0.125, 0.25], [0.25, 0.5], [0.5, 1.0]] 31 | regress_range = [0.] + [1 / (2 ** (index - 1)) for index in range(len(points), 0, -1)] 32 | regress_range = [[regress_range[index], regress_range[index + 1]] for index in range(len(points))] 33 | regress_range = [tf.ones_like(p) * r for p, r in zip(points, regress_range)] 34 | 35 | args = [arg for arg in [y_pred, bbox_pred, points, regress_range, conf_pred] if arg is not None] 36 | out = AnchorFreeLoss(class_loss = class_loss, bbox_loss = bbox_loss, conf_loss = conf_loss, 37 | decode_bbox = decode_bbox, weight = class_weight, background = background, 38 | assign = assign, sampler = sampler, 39 | batch_size = batch_size, 40 | missing_value = missing_value, dtype = tf.float32, name = "anchor_free_loss")([y_true, bbox_true], args) 41 | args = [arg for arg in [y_pred, bbox_pred, points, conf_pred] if arg is not None] 42 | y_pred, bbox_pred = FilterDetection(proposal_count = proposal_count, iou_threshold = iou_threshold, score_threshold = score_threshold, soft_nms = soft_nms, ignore_label = ignore_label, performance_count = performance_count, 43 | batch_size = batch_size, dtype = tf.float32, name = "filter_detection")(args) 44 | model = tf.keras.Model([input, y_true, bbox_true], [y_pred, bbox_pred]) 45 | 46 | loss_class, loss_bbox = out[:2] 47 | loss_conf = out[2] if 2 < len(out) else None 48 | loss_class = [loss_class] if not isinstance(loss_class, (tuple, list)) else loss_class 49 | loss_bbox = [loss_bbox] if not isinstance(loss_bbox, (tuple, list)) else loss_bbox 50 | loss_conf = [loss_conf] if loss_conf is not None and not isinstance(loss_conf, (tuple, list)) else loss_conf 51 | #for level, (_loss_class, _loss_bbox) in enumerate(zip(loss_class, loss_bbox)): 52 | # model.add_loss(tf.expand_dims(_loss_class, axis = -1)) 53 | # model.add_loss(tf.expand_dims(_loss_bbox, axis = -1)) 54 | # if loss_conf is not None: 55 | # model.add_loss(tf.expand_dims(loss_conf[level], axis = -1)) 56 | 57 | losses = [] 58 | loss_class = tf.expand_dims(tf.reduce_sum(loss_class), axis = -1) 59 | loss_bbox = tf.expand_dims(tf.reduce_sum(loss_bbox), axis = -1) 60 | model.add_metric(loss_class, name = "loss_class", aggregation = "mean") 61 | model.add_metric(loss_bbox, name = "loss_bbox", aggregation = "mean") 62 | losses += [loss_class, loss_bbox] 63 | if loss_conf is not None: 64 | loss_conf = tf.expand_dims(tf.reduce_sum(loss_conf), axis = -1) 65 | model.add_metric(loss_conf, name = "loss_conf", aggregation = "mean") 66 | losses += [loss_conf] 67 | 68 | losses = tf.reduce_sum(losses, axis = 0) 69 | model.add_loss(losses) 70 | 71 | if regularize: 72 | model.add_loss(lambda: tf.cast(tf.reduce_sum(regularize_loss(model, weight_decay), keepdims = True), tf.float32)) 73 | return model -------------------------------------------------------------------------------- /tfdet/model/train/loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .anchor import AnchorLoss 2 | from .anchor_free import AnchorFreeLoss 3 | from .roi import RoiTarget, RoiBboxLoss, RoiMaskLoss 4 | from .mask import FusedSemanticLoss, ResizeMaskLoss 5 | from .yolo import YoloLoss 6 | from . import util -------------------------------------------------------------------------------- /tfdet/model/train/loss/mask.py: -------------------------------------------------------------------------------- 1 | import functools 2 | 3 | import tensorflow as tf 4 | 5 | from tfdet.core.loss import categorical_cross_entropy, resize_loss 6 | 7 | class FusedSemanticLoss(tf.keras.layers.Layer): 8 | def __init__(self, loss = categorical_cross_entropy, 9 | weight = None, method = "bilinear", missing_value = 0., dtype = tf.float32, **kwargs): 10 | kwargs["dtype"] = dtype 11 | super(FusedSemanticLoss, self).__init__(**kwargs) 12 | self._loss = loss 13 | self.weight = weight 14 | self.method = method 15 | self.missing_value = missing_value 16 | 17 | _loss = functools.partial(self.loss, loss = self._loss, weight = self.weight, method = self.method, missing_value = self.missing_value) 18 | self.loss_func = tf.keras.layers.Lambda(lambda args: _loss(*args), dtype = self.dtype, name = "loss") 19 | 20 | @staticmethod 21 | @tf.function 22 | def loss(y_true, mask_true, semantic_pred, loss = categorical_cross_entropy, weight = None, method = "bilinear", missing_value = 0.): 23 | """ 24 | Args: 25 | y_true = #(N, padded_num_true, 1 or num_class) 26 | mask_true = #(N, padded_num_true, h, w, 1) 27 | semantic_pred = #(N, h, w, n_class) 28 | 29 | Returns: 30 | semantic_loss = #(1,) 31 | """ 32 | mask_shape = tf.shape(mask_true) 33 | semantic_shape = tf.shape(semantic_pred) 34 | 35 | label_true = tf.cond(tf.equal(tf.shape(y_true)[2], 1), true_fn = lambda: tf.cast(y_true[..., 0], tf.int32), false_fn = lambda: tf.argmax(y_true, axis = -1, output_type = tf.int32)) 36 | 37 | semantic_true = tf.reshape(mask_true, [-1, mask_shape[2], mask_shape[3], 1]) 38 | semantic_true = tf.image.resize(semantic_true, semantic_shape[-3:-1], method = method) 39 | semantic_true = tf.reshape(semantic_true, [mask_shape[0], mask_shape[1], semantic_shape[-3], semantic_shape[-2]]) 40 | semantic_true = tf.clip_by_value(tf.round(semantic_true), 0., 1.) 41 | semantic_true = tf.multiply(tf.expand_dims(tf.expand_dims(label_true, axis = -1), axis = -1), tf.cast(semantic_true, tf.int32)) 42 | semantic_true = tf.reduce_max(semantic_true, axis = 1) 43 | semantic_true = tf.one_hot(semantic_true, semantic_shape[-1]) 44 | 45 | _loss = loss(semantic_true, semantic_pred, weight = weight) 46 | _loss = tf.where(tf.logical_or(tf.math.is_nan(_loss), tf.math.is_inf(_loss)), tf.cast(missing_value, _loss.dtype), _loss) 47 | return _loss 48 | 49 | def call(self, inputs, outputs, bbox = True, mask = True): 50 | y_true, mask_true = inputs[:2] 51 | if 2 < len(inputs) and tf.keras.backend.int_shape(mask_true)[-1] == 4: 52 | mask_true = inputs[2] 53 | semantic_pred = outputs[-1] if isinstance(outputs, (tuple, list)) else outputs 54 | 55 | semantic_loss = self.loss_func([y_true, mask_true, semantic_pred]) 56 | return semantic_loss 57 | 58 | class ResizeMaskLoss(tf.keras.layers.Layer): 59 | def __init__(self, loss = categorical_cross_entropy, 60 | weight = None, method = "bilinear", missing_value = 0., dtype = tf.float32, **kwargs): 61 | kwargs["dtype"] = dtype 62 | super(ResizeMaskLoss, self).__init__(**kwargs) 63 | self._loss = loss 64 | self.weight = weight 65 | self.method = method 66 | self.missing_value = missing_value 67 | 68 | _loss = functools.partial(self.loss, loss = self._loss, weight = self.weight, method = self.method, missing_value = self.missing_value) 69 | self.loss_func = tf.keras.layers.Lambda(lambda args: _loss(*args), dtype = self.dtype, name = "loss") 70 | 71 | @staticmethod 72 | @tf.function 73 | def loss(mask_true, mask_pred, loss = categorical_cross_entropy, weight = None, method = "bilinear", missing_value = 0.): 74 | """ 75 | Args: 76 | mask_true = #(N, h, w, 1 or n_class) 77 | mask_pred = #(N, h, w, n_class) 78 | 79 | Returns: 80 | semantic_loss = #(1,) 81 | """ 82 | _loss = resize_loss(mask_true, mask_pred, loss = loss, method = method, weight = weight, reduce = None) 83 | _loss = tf.reduce_mean(_loss) 84 | _loss = tf.where(tf.logical_or(tf.math.is_nan(_loss), tf.math.is_inf(_loss)), tf.cast(missing_value, _loss.dtype), _loss) 85 | return _loss 86 | 87 | def call(self, inputs, outputs): 88 | mask_true = inputs[-1] if isinstance(inputs, (tuple, list)) else inputs 89 | mask_pred = outputs[-1] if isinstance(outputs, (tuple, list)) else outputs 90 | loss = self.loss_func([mask_true, mask_pred]) 91 | return loss -------------------------------------------------------------------------------- /tfdet/model/train/loss/util.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def image_to_level(tensor, n_level): 4 | out = [] 5 | start = 0 6 | for n in n_level: 7 | end = start + n 8 | out.append(tensor[:, start:end]) 9 | start = end 10 | return out -------------------------------------------------------------------------------- /tfdet/model/train/padim.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from ..head.padim import FeatureExtractor 5 | 6 | def decode(fv): 7 | w = np.identity(np.shape(fv)[-1]) * 0.01 8 | cvar = np.array([np.cov(f, rowvar = False) for f in fv]) + w 9 | cvar_inv = np.array([np.linalg.inv(v) for v in cvar]) 10 | return cvar_inv 11 | 12 | def train(feature): 13 | if tf.is_tensor(feature): 14 | b, h, w, c = tf.keras.backend.int_shape(feature) 15 | feature = tf.reshape(feature, [-1, h * w, c]) 16 | mean = tf.reduce_mean(feature, axis = 0) 17 | cvar_inv = tf.py_function(lambda *args: decode(args[0].numpy()), inp = [tf.transpose(feature, [1, 0, 2])], Tout = feature.dtype) 18 | cvar_inv = tf.reshape(cvar_inv, [h * w, c, c]) 19 | else: 20 | b, h, w, c = np.shape(feature) 21 | feature = np.reshape(feature, [-1, h * w, c]) 22 | mean = np.mean(feature, axis = 0) 23 | cvar_inv = decode(np.transpose(feature, [1, 0, 2])) 24 | cvar_inv = np.reshape(cvar_inv, [h * w, c, c]) 25 | return mean, cvar_inv -------------------------------------------------------------------------------- /tfdet/model/train/patch_core.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from tfdet.core.ops import core_sampling 5 | from ..head.patch_core import FeatureExtractor 6 | 7 | def train(feature, n_sample = 0.001, n_feature = "auto", eps = 0.9): 8 | if tf.is_tensor(feature): 9 | b, h, w, c = tf.keras.backend.int_shape(feature) 10 | feature = tf.reshape(feature, [-1, c]) 11 | feature = tf.py_function(lambda *args: core_sampling(args[0].numpy(), n_sample = n_sample, n_feature = n_feature, eps = eps), inp = [feature], Tout = feature.dtype) 12 | feature = tf.reshape(feature, [-1, c]) 13 | else: 14 | b, h, w, c = np.shape(feature) 15 | feature = np.reshape(feature, [-1, c]) 16 | feature = core_sampling(feature, n_sample = n_sample, n_feature = n_feature, eps = eps) 17 | feature = np.reshape(feature, [-1, c]) 18 | return feature -------------------------------------------------------------------------------- /tfdet/model/train/retina.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.assign import max_iou 4 | from tfdet.core.loss import focal_binary_cross_entropy, smooth_l1, regularize as regularize_loss 5 | from .loss import AnchorLoss 6 | from ..postprocess.anchor import FilterDetection 7 | 8 | def train_model(input, y_pred, bbox_pred, anchors, 9 | assign = max_iou, sampler = None, valid_inside_anchor = False, 10 | proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, ignore_label = 0, performance_count = 5000, 11 | mean = [0., 0., 0., 0.], std = [1., 1., 1., 1.], clip_ratio = 16 / 1000, 12 | class_loss = focal_binary_cross_entropy, bbox_loss = smooth_l1, 13 | regularize = True, weight_decay = 1e-4, 14 | decode_bbox = False, class_weight = None, background = False, 15 | batch_size = 1, missing_value = 0.): 16 | y_true = tf.keras.layers.Input(shape = (None, None), name = "y_true") 17 | bbox_true = tf.keras.layers.Input(shape = (None, 4), name = "bbox_true") 18 | 19 | loss_class, loss_bbox = AnchorLoss(class_loss = class_loss, bbox_loss = bbox_loss, 20 | decode_bbox = decode_bbox, valid_inside_anchor = valid_inside_anchor, weight = class_weight, background = background, 21 | assign = assign, sampler = sampler, 22 | mean = mean, std = std, clip_ratio = clip_ratio, 23 | batch_size = batch_size, 24 | missing_value = missing_value, dtype = tf.float32, name = "anchor_loss")([y_true, bbox_true], [y_pred, bbox_pred, anchors]) 25 | y_pred, bbox_pred = FilterDetection(proposal_count = proposal_count, iou_threshold = iou_threshold, score_threshold = score_threshold, soft_nms = soft_nms, valid_inside_anchor = valid_inside_anchor, ignore_label = ignore_label, performance_count = performance_count, 26 | mean = mean, std = std, clip_ratio = clip_ratio, 27 | batch_size = batch_size, dtype = tf.float32, name = "filter_detection")([y_pred, bbox_pred, anchors]) 28 | model = tf.keras.Model([input, y_true, bbox_true], [y_pred, bbox_pred]) 29 | 30 | loss_class = [loss_class] if not isinstance(loss_class, (tuple, list)) else loss_class 31 | loss_bbox = [loss_bbox] if not isinstance(loss_bbox, (tuple, list)) else loss_bbox 32 | #for _loss_class, _loss_bbox in zip(loss_class, loss_bbox): 33 | # model.add_loss(tf.expand_dims(_loss_class, axis = -1)) 34 | # model.add_loss(tf.expand_dims(_loss_bbox, axis = -1)) 35 | 36 | loss_class = tf.expand_dims(tf.reduce_sum(loss_class), axis = -1) 37 | loss_bbox = tf.expand_dims(tf.reduce_sum(loss_bbox), axis = -1) 38 | model.add_metric(loss_class, name = "loss_class", aggregation = "mean") 39 | model.add_metric(loss_bbox, name = "loss_bbox", aggregation = "mean") 40 | 41 | losses = tf.reduce_sum([loss_class, loss_bbox], axis = 0) 42 | model.add_loss(losses) 43 | 44 | if regularize: 45 | model.add_loss(lambda: tf.cast(tf.reduce_sum(regularize_loss(model, weight_decay), keepdims = True), tf.float32)) 46 | return model -------------------------------------------------------------------------------- /tfdet/model/train/segmentation.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.loss import categorical_cross_entropy, regularize as regularize_loss 4 | from .loss import ResizeMaskLoss 5 | 6 | def train_model(input, mask_pred, aux_mask_pred = None, 7 | loss = categorical_cross_entropy, 8 | regularize = True, weight_decay = 1e-4, 9 | class_weight = None, method = "bilinear", aux_weight = 0.4, 10 | missing_value = 0.): 11 | mask_pred = [mask_pred] if not isinstance(mask_pred, (tuple, list)) else mask_pred 12 | if aux_mask_pred is None and 1 < len(mask_pred): 13 | aux_mask_pred = mask_pred[1] 14 | mask_pred = mask_pred[0] 15 | 16 | mask_true = tf.keras.layers.Input(shape = (None, None, None), dtype = tf.uint16, name = "mask_true") 17 | _loss = ResizeMaskLoss(loss = loss, weight = class_weight, method = method, missing_value = missing_value, dtype = tf.float32, name = "mask_loss")(mask_true, mask_pred) 18 | aux_loss = None 19 | if aux_mask_pred is not None: 20 | aux_loss = ResizeMaskLoss(loss = loss, weight = class_weight, method = method, missing_value = missing_value, dtype = tf.float32, name = "aux_mask_loss")(mask_true, aux_mask_pred) 21 | aux_loss *= aux_weight 22 | model = tf.keras.Model([input, mask_true], mask_pred) 23 | 24 | _loss = tf.expand_dims(_loss, axis = -1) 25 | if aux_loss is not None: 26 | model.add_metric(_loss, name = "loss_mask", aggregation = "mean") 27 | aux_loss = tf.expand_dims(aux_loss, axis = -1) 28 | model.add_metric(aux_loss, name = "loss_aux_mask", aggregation = "mean") 29 | _loss = tf.reduce_sum([_loss, aux_loss], axis = 0) 30 | model.add_metric(_loss, name = "loss", aggregation = "mean") 31 | model.add_loss(_loss) 32 | 33 | if regularize: 34 | model.add_loss(lambda: tf.cast(tf.reduce_sum(regularize_loss(model, weight_decay), keepdims = True), tf.float32)) 35 | return model -------------------------------------------------------------------------------- /tfdet/model/train/spade.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from ..head.spade import FeatureExtractor 4 | 5 | def train(feature): 6 | return feature 7 | -------------------------------------------------------------------------------- /tfdet/model/train/yolo.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from tfdet.core.assign import max_iou 4 | from tfdet.core.loss import binary_cross_entropy, focal_binary_cross_entropy, ciou, regularize as regularize_loss 5 | from .loss import YoloLoss 6 | from ..postprocess.yolo import FilterDetection 7 | 8 | def focal_loss(y_true, y_pred, alpha = .25, gamma = 2., weight = None, reduce = tf.reduce_mean): 9 | return focal_binary_cross_entropy(y_true, y_pred, alpha = alpha, gamma = gamma, weight = weight, reduce = reduce) 10 | 11 | def train_model(input, score_pred, logit_pred, bbox_pred, anchors, 12 | assign = max_iou, sampler = None, 13 | proposal_count = 100, iou_threshold = 0.5, score_threshold = 0.05, soft_nms = False, valid_inside_anchor = False, ignore_label = 0, performance_count = 5000, 14 | clip_ratio = 16 / 1000, 15 | score_loss = binary_cross_entropy, class_loss = focal_loss, bbox_loss = ciou, 16 | regularize = True, weight_decay = 1e-4, 17 | decode_bbox = True, class_weight = None, 18 | batch_size = 1, missing_value = 0.): 19 | y_true = tf.keras.layers.Input(shape = (None, None), name = "y_true") 20 | bbox_true = tf.keras.layers.Input(shape = (None, 4), name = "bbox_true") 21 | 22 | loss_score, loss_class, loss_bbox = YoloLoss(score_loss = score_loss, class_loss = class_loss, bbox_loss = bbox_loss, 23 | decode_bbox = decode_bbox, valid_inside_anchor = valid_inside_anchor, weight = class_weight, 24 | assign = assign, sampler = sampler, 25 | clip_ratio = clip_ratio, 26 | batch_size = batch_size, 27 | missing_value = missing_value, dtype = tf.float32, name = "yolo_loss")([y_true, bbox_true], [score_pred, logit_pred, bbox_pred, anchors]) 28 | y_pred, bbox_pred = FilterDetection(proposal_count = proposal_count, iou_threshold = iou_threshold, score_threshold = score_threshold, soft_nms = soft_nms, valid_inside_anchor = valid_inside_anchor, ignore_label = ignore_label, performance_count = performance_count, 29 | clip_ratio = clip_ratio, 30 | batch_size = batch_size, dtype = tf.float32, name = "filter_detection")([score_pred, logit_pred, bbox_pred, anchors]) 31 | model = tf.keras.Model([input, y_true, bbox_true], [y_pred, bbox_pred]) 32 | 33 | loss_score = [loss_score] if not isinstance(loss_score, (tuple, list)) else loss_score 34 | loss_class = [loss_class] if not isinstance(loss_class, (tuple, list)) else loss_class 35 | loss_bbox = [loss_bbox] if not isinstance(loss_bbox, (tuple, list)) else loss_bbox 36 | #for _loss_score, _loss_class, _loss_bbox in zip(loss_score, loss_class, loss_bbox): 37 | # model.add_loss(tf.expand_dims(_loss_score, axis = -1)) 38 | # model.add_loss(tf.expand_dims(_loss_class, axis = -1)) 39 | # model.add_loss(tf.expand_dims(_loss_bbox, axis = -1)) 40 | 41 | loss_score = tf.expand_dims(tf.reduce_sum(loss_score), axis = -1) 42 | loss_class = tf.expand_dims(tf.reduce_sum(loss_class), axis = -1) 43 | loss_bbox = tf.expand_dims(tf.reduce_sum(loss_bbox), axis = -1) 44 | model.add_metric(loss_score, name = "loss_score", aggregation = "mean") 45 | model.add_metric(loss_class, name = "loss_class", aggregation = "mean") 46 | model.add_metric(loss_bbox, name = "loss_bbox", aggregation = "mean") 47 | 48 | losses = tf.reduce_sum([loss_score, loss_class, loss_bbox], axis = 0) 49 | model.add_loss(losses) 50 | 51 | if regularize: 52 | model.add_loss(lambda: tf.cast(tf.reduce_sum(regularize_loss(model, weight_decay), keepdims = True), tf.float32)) 53 | return model -------------------------------------------------------------------------------- /tfdet/util/__init__.py: -------------------------------------------------------------------------------- 1 | from .metric import get_threshold 2 | from .visualize import draw_bbox -------------------------------------------------------------------------------- /tfdet/util/metric.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | def get_threshold(y_true, y_pred): 5 | try: 6 | from sklearn.metrics import roc_auc_score, roc_curve, precision_recall_curve 7 | except Exception as e: 8 | print("If you want to use 'get_threshold', please install 'scikit-learn 0.14▲'") 9 | raise e 10 | precision, recall, thresholds = precision_recall_curve(y_true, y_pred) 11 | f1 = np.divide(2 * precision * recall, precision + recall, out = np.zeros_like(precision), where = (precision + recall) != 0) 12 | threshold = thresholds[np.argmax(f1)] 13 | return threshold 14 | -------------------------------------------------------------------------------- /tfdet/util/visualize.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | 4 | def draw_bbox(x_true, bbox_true, y_true = None, mask_true = None, label = None, crop_mask = True, threshold = 0.5, mix_ratio = 0.5, method = cv2.INTER_LINEAR, probability = True, prefix = "", postfix = "", color = None, size_ratio = 1.): 5 | batch = True 6 | if not isinstance(x_true, (tuple, list)) and np.ndim(x_true) not in [1, 4]: 7 | batch = False 8 | x_true = [x_true] 9 | bbox_true = [bbox_true] 10 | if y_true is not None: 11 | y_true = [y_true] 12 | if mask_true is not None: 13 | mask_true = [mask_true] 14 | 15 | result = [] 16 | for batch_index in range(len(x_true)): 17 | image = np.array(x_true[batch_index]) 18 | bbox = bbox_true[batch_index] 19 | bbox = np.array(bbox) if not isinstance(bbox, np.ndarray) else bbox 20 | h, w, ch = np.shape(image)[:3] 21 | size = int(max(h, w) / 500 * size_ratio) 22 | font_size = max(h, w) / 1250 * size_ratio 23 | #normalize_flag = np.max(image) <= 1 24 | normalize_flag = image.dtype != np.uint8 25 | y_color = (1, 1, 1) if normalize_flag else (255, 255, 255) 26 | if ch == 1: 27 | y_color = y_color[0] 28 | 29 | #valid_indices = np.where(0 < np.max(bbox, axis = -1))[0] 30 | valid_indices = np.where(np.any(0 < bbox, axis = -1))[0] 31 | bbox = bbox[valid_indices] 32 | mask = mask_true[batch_index] if mask_true is not None else None 33 | mask = (np.array(mask) if not isinstance(mask, np.ndarray) else mask)[valid_indices] if mask is not None else None 34 | if y_true is not None: 35 | y = y_true[batch_index] 36 | y = (np.array(y) if not isinstance(y, np.ndarray) else y)[valid_indices] 37 | if np.shape(y)[-1] != 1: 38 | y_index = np.argmax(y, axis = -1) 39 | score = np.max(y, axis = -1) 40 | else: 41 | y_index = y[..., 0] 42 | score = np.ones_like(y_index, dtype = np.float32) 43 | 44 | for index, rect in enumerate(bbox): 45 | bbox_color = color 46 | if color is None or (y_true is not None and not np.issubdtype(y_index.dtype, np.number) and not (label and np.ndim(bbox_color) == 2)): 47 | bbox_color = np.random.random(size = ch) if normalize_flag else np.random.randint(0, 256, size = ch).astype(float) 48 | if np.max(rect) < 2: 49 | rect = np.round(np.multiply(rect, [w, h, w, h])) 50 | rect = tuple(rect.astype(int)) 51 | 52 | if y_true is not None: 53 | if not np.issubdtype(y_index.dtype, np.number): 54 | name = y_index[index] 55 | bbox_color = bbox_color[np.argmax(np.isin(label, name))] if np.ndim(bbox_color) == 2 and label is not None else bbox_color 56 | else: 57 | name = label[int(y_index[index])] if label is not None else int(y_index[index]) 58 | bbox_color = bbox_color[y_index[index]] if np.ndim(bbox_color) == 2 else bbox_color 59 | msg = "{0}{1}".format(prefix, name) 60 | if probability: 61 | msg = "{0}:{1:.2f}".format(msg, score[index]) 62 | msg = "{0}{1}".format(msg, postfix) 63 | text_size = cv2.getTextSize(msg, cv2.FONT_HERSHEY_SIMPLEX, font_size, size)[0] 64 | font_pos = (rect[0], max(rect[1], text_size[1])) 65 | 66 | bbox_color = tuple(bbox_color) if np.ndim(bbox_color) == 1 else bbox_color 67 | cv2.rectangle(image, rect[:2], rect[-2:], bbox_color, size) 68 | if y_true is not None: 69 | cv2.rectangle(image, (font_pos[0], font_pos[1] - text_size[1]), (font_pos[0] + text_size[0], font_pos[1]), bbox_color, -1) 70 | cv2.putText(image, msg, font_pos, cv2.FONT_HERSHEY_SIMPLEX, font_size, y_color, size) 71 | 72 | if mask_true is not None: 73 | m = mask[index] 74 | mh, mw = np.shape(m)[:2] 75 | if crop_mask: 76 | mask_rect = np.multiply(np.divide(rect, [w, h, w, h]), [mw, mh, mw, mh]).astype(np.int32) 77 | m = m[mask_rect[1]:mask_rect[3], mask_rect[0]:mask_rect[2]] 78 | crop = image[rect[1]:rect[3], rect[0]:rect[2]] 79 | m = cv2.resize(m, np.shape(crop)[:2][::-1], interpolation = method) 80 | m = np.where(np.greater(m, threshold), 1., 0.) 81 | m = np.tile(np.expand_dims(m, axis = -1), [1, 1, ch]) * bbox_color 82 | image[rect[1]:rect[3], rect[0]:rect[2]] = np.where(np.greater(m, 0), crop * (1 - mix_ratio) + m * mix_ratio, crop) 83 | result.append(image) 84 | if not batch:# and len(x_true) == 1: 85 | result = result[0] 86 | else: 87 | try: 88 | result = np.array(result) 89 | except: 90 | pass 91 | return result 92 | --------------------------------------------------------------------------------