Merge pull request #96 from BloodAxe/feature/modules-overhaul

Feature/modules overhaul

.github/workflows/ci.yml

@@ -51,6 +51,6 @@ jobs:
     - name: Update pip
       run: python -m pip install --upgrade pip
     - name: Install Black
-      run: pip install black==23.3.0
+      run: pip install black==23.7.0
     - name: Run Black
       run: black --config=pyproject.toml --check .

.github/workflows/codeql.yml

@@ -0,0 +1,66 @@
+name: "CodeQL"
+
+on:
+  push:
+    branches: [ 'master' ]
+  pull_request:
+    # The branches below must be a subset of the branches above
+    branches: [ 'master' ]
+  schedule:
+    - cron: '39 12 * * 5'
+
+jobs:
+  analyze:
+    name: Analyze
+    runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-latest' }}
+    timeout-minutes: ${{ (matrix.language == 'swift' && 120) || 360 }}
+    permissions:
+      actions: read
+      contents: read
+      security-events: write
+
+    strategy:
+      fail-fast: false
+      matrix:
+        language: [ 'python' ]
+        # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python', 'ruby' ]
+        # Use only 'java' to analyze code written in Java, Kotlin or both
+        # Use only 'javascript' to analyze code written in JavaScript, TypeScript or both
+        # Learn more about CodeQL language support at https://aka.ms/codeql-docs/language-support
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v3
+
+    # Initializes the CodeQL tools for scanning.
+    - name: Initialize CodeQL
+      uses: github/codeql-action/init@v2
+      with:
+        languages: ${{ matrix.language }}
+        # If you wish to specify custom queries, you can do so here or in a config file.
+        # By default, queries listed here will override any specified in a config file.
+        # Prefix the list here with "+" to use these queries and those in the config file.
+
+        # For more details on CodeQL's query packs, refer to: https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs
+        # queries: security-extended,security-and-quality
+
+
+    # Autobuild attempts to build any compiled languages (C/C++, C#, Go, Java, or Swift).
+    # If this step fails, then you should remove it and run the build manually (see below)
+    - name: Autobuild
+      uses: github/codeql-action/autobuild@v2
+
+    # ℹ️ Command-line programs to run using the OS shell.
+    # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun
+
+    #   If the Autobuild fails above, remove it and uncomment the following three lines.
+    #   modify them (or add more) to build your code if your project, please refer to the EXAMPLE below for guidance.
+
+    # - run: |
+    #     echo "Run, Build Application using script"
+    #     ./location_of_script_within_repo/buildscript.sh
+
+    - name: Perform CodeQL Analysis
+      uses: github/codeql-action/analyze@v2
+      with:
+        category: "/language:${{matrix.language}}"

makefile

@@ -8,5 +8,8 @@ publish: build
 	python -m twine check dist/*
 	python -m twine upload dist/*
 
+black:
+	black --config=pyproject.toml .
+
 clean:
 	rm -r build dist *.egg-info || true

pytorch_toolbelt/datasets/providers/__init__.py

@@ -1 +1,2 @@
 from .inria_aerial import *
+from .coco_detection import *

pytorch_toolbelt/datasets/providers/coco_detection.py

@@ -0,0 +1,112 @@
+import dataclasses
+import json
+import os.path
+from collections import defaultdict
+from typing import List
+
+import numpy as np
+from torch.utils.data import Dataset
+
+__all__ = ["DetectionSample", "COCODetectionDatasetReader"]
+
+
+@dataclasses.dataclass
+class DetectionSample:
+    image_id: str
+    image_path: str
+    image_width: int
+    image_height: int
+
+    labels: np.ndarray  # [N]
+    bboxes: np.ndarray  # [N, 4] in XYXY format
+    is_difficult: np.ndarray  # [N]
+
+
+class COCODetectionDatasetReader(Dataset):
+    samples: List[DetectionSample]
+    class_names: List[str]
+    num_classes: int
+
+    def __init__(self, samples: List[DetectionSample], class_names: List[str]):
+        self.samples = samples
+        self.class_names = class_names
+        self.num_classes = len(class_names)
+
+    def __getitem__(self, item):
+        return self.samples[item]
+
+    def __len__(self):
+        return len(self.samples)
+
+    @staticmethod
+    def convert_to_dict(annotations):
+        result_dict = defaultdict(list)
+        for obj in annotations:
+            image_id = obj["image_id"]
+            result_dict[image_id].append(obj)
+        return result_dict
+
+    @classmethod
+    def from_directory_and_annotation(cls, images_directory: str, annotation: str):
+        samples = []
+        with open(annotation, "r") as f:
+            data = json.load(f)
+
+        category_ids, class_names = zip(*[(category["id"], category["name"]) for category in data["categories"]])
+
+        annotations = cls.convert_to_dict(data["annotations"])
+        category_id_to_index = {category_id: index for index, category_id in enumerate(category_ids)}
+
+        for image in data["images"]:
+            image_id = image["id"]
+            image_path = os.path.join(images_directory, image["file_name"])
+            image_width = image["width"]
+            image_height = image["height"]
+
+            labels = []
+            bboxes = []
+            is_difficult = []
+
+            if image_id in annotations:
+                for annotations in annotations[image_id]:
+                    class_index = category_id_to_index[annotations["category_id"]]
+                    x, y, w, h = annotations["bbox"]
+                    bbox_xyxy = [x, y, x + w, y + h]
+
+                    labels.append(class_index)
+                    bboxes.append(bbox_xyxy)
+                    is_difficult.append(annotations["iscrowd"])
+
+            sample = DetectionSample(
+                image_id=image_id,
+                image_path=image_path,
+                image_width=image_width,
+                image_height=image_height,
+                labels=np.array(labels, dtype=int).reshape(-1),
+                bboxes=np.array(bboxes, dtype=np.float32).reshape(-1, 4),
+                is_difficult=np.array(is_difficult, dtype=bool).reshape(-1),
+            )
+            samples.append(sample)
+
+        return cls(samples, class_names)
+
+
+if __name__ == "__main__":
+    import cv2
+
+    start = cv2.getTickCount()
+    train_ds = COCODetectionDatasetReader.from_directory_and_annotation(
+        images_directory="e:/coco2017/images/train2017/", annotation="e:/coco2017/annotations/instances_train2017.json"
+    )
+    end = cv2.getTickCount()
+    print((end - start) / cv2.getTickFrequency())
+
+    start = cv2.getTickCount()
+    valid_ds = COCODetectionDatasetReader.from_directory_and_annotation(
+        images_directory="e:/coco2017/images/val2017/", annotation="e:/coco2017/annotations/instances_val2017.json"
+    )
+    end = cv2.getTickCount()
+    print((end - start) / cv2.getTickFrequency())
+
+    print(len(train_ds))
+    print(len(valid_ds))

pytorch_toolbelt/datasets/wrappers.py

@@ -1,4 +1,5 @@
 import random
+import typing
 from typing import Any, Optional
 
 from torch.utils.data import Dataset
@@ -74,8 +75,8 @@ def __getitem__(self, _) -> Any:
         index = random.choice(self.indexes)
         return self.dataset[index]
 
-    def get_collate_fn(self):
+    def get_collate_fn(self) -> typing.Callable:
         get_collate_fn = getattr(self.dataset, "get_collate_fn", None)
         if callable(get_collate_fn):
             return get_collate_fn()
-        return default_collate()
+        return default_collate

pytorch_toolbelt/inference/functional.py

@@ -1,3 +1,4 @@
+import itertools
 from collections.abc import Sized, Iterable
 from typing import Union, Tuple
 
@@ -13,6 +14,7 @@
     "logodd_mean",
     "log1p_mean",
     "pad_image_tensor",
+    "pad_tensor_to_size",
     "torch_fliplr",
     "torch_flipud",
     "torch_none",
@@ -143,6 +145,37 @@ def torch_transpose2(x: Tensor):
     return x.transpose(3, 2)
 
 
+def pad_tensor_to_size(x: Tensor, size: Tuple[int, ...], mode="constant", value=0) -> Tuple[Tensor, Tuple[slice, ...]]:
+    """
+    Pad tensor to given size by appending elements to the beginning and end of each axis.
+
+    :param x: Input tensor of shape [B, C, *num_spatial_dims]
+    :param size: Target tensor size defined as an array of [num_spatial_dims] elements
+    :param mode: Padding mode, see torch.nn.functional.pad
+    :param value: Padding value, see torch.nn.functional.pad
+    :return: Tuple of padded tensor and crop parameters. Second argument can be used to reverse pad operation of model output
+    """
+    num_spatial_dims = len(size)
+    if num_spatial_dims != len(x.shape) - 2:
+        raise ValueError(f"Expected {num_spatial_dims} spatial dimensions, got {len(x.shape) - 2}")
+
+    spatial_dims = x.shape[-num_spatial_dims:]
+    padding = torch.tensor(size) - torch.tensor(spatial_dims)
+    padding_before = padding // 2
+    padding_after = padding - padding_before
+
+    padding_pairs = tuple(zip(padding_before.tolist(), padding_after.tolist()))
+    padding_params = tuple(itertools.chain(*reversed(padding_pairs)))
+
+    x = torch.nn.functional.pad(x, pad=padding_params, mode=mode, value=value)
+
+    crop_params = [slice(None), slice(None)] + [
+        slice(before, before + total_size)
+        for (before, after, total_size) in zip(padding_before, padding_after, spatial_dims)
+    ]
+    return x, crop_params
+
+
 def pad_image_tensor(
     image_tensor: Tensor, pad_size: Union[int, Tuple[int, int]] = 32
 ) -> Tuple[Tensor, Tuple[int, int, int, int]]:

pytorch_toolbelt/inference/tta.py

@@ -677,7 +677,7 @@ def ms_image_deaugment(
         if rows_offset == 0 and cols_offset == 0:
             deaugmented_outputs.append(feature_map)
         else:
-            batch_size, channels, rows, cols = feature_map.size()
+            _, _, rows, cols = feature_map.size()
             original_size = rows - rows_offset // stride, cols - cols_offset // stride
             scaled_image = torch.nn.functional.interpolate(
                 feature_map, size=original_size, mode=mode, align_corners=align_corners
@@ -760,6 +760,9 @@ def __init__(
         self,
         model: nn.Module,
         size_offsets: List[int],
+        mode: str = "bilinear",
+        align_corners: bool = False,
+        augment_fn: Callable = ms_image_augment,
         deaugment_fn: Union[Callable, Dict[str, Callable]] = ms_image_deaugment,
     ):
         if isinstance(deaugment_fn, Mapping):
@@ -770,10 +773,15 @@ def __init__(
         super().__init__()
         self.model = model
         self.size_offsets = size_offsets
+        self.mode = mode
+        self.align_corners = align_corners
+        self.augment_fn = augment_fn
         self.deaugment_fn = deaugment_fn
 
     def forward(self, x):
-        ms_inputs = ms_image_augment(x, size_offsets=self.size_offsets)
+        ms_inputs = self.augment_fn(
+            x, size_offsets=self.size_offsets, mode=self.mode, align_corners=self.align_corners
+        )
         ms_outputs = [self.model(x) for x in ms_inputs]
 
         outputs = {}
@@ -785,6 +793,8 @@ def forward(self, x):
             for key in keys:
                 deaugment_fn: Callable = self.deaugment_fn[key]
                 values = [x[key] for x in ms_outputs]
-                outputs[key] = deaugment_fn(values, self.size_offsets)
+                outputs[key] = deaugment_fn(
+                    values, size_offsets=self.size_offsets, mode=self.mode, align_corners=self.align_corners
+                )
 
         return outputs

pytorch_toolbelt/losses/focal.py

@@ -11,6 +11,8 @@
 
 
 class BinaryFocalLoss(nn.Module):
+    __constants__ = ["alpha", "gamma", "reduction", "ignore_index", "normalized", "reduced_threshold", "activation"]
+
     def __init__(
         self,
         alpha: Optional[float] = None,
@@ -33,7 +35,15 @@ def __init__(
 
         """
         super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
         self.ignore_index = ignore_index
+        self.reduction = reduction
+        self.normalized = normalized
+        self.reduced_threshold = reduced_threshold
+        self.activation = activation
+        self.softmax_dim = softmax_dim
+
         self.focal_loss_fn = partial(
             focal_loss_with_logits,
             alpha=alpha,
@@ -50,6 +60,13 @@ def __init__(
             self._one_hot_targets_with_ignore if ignore_index is not None else self._one_hot_targets
         )
 
+    def __repr__(self):
+        repr = f"{self.__class__.__name__}(alpha={self.alpha}, gamma={self.gamma}, "
+        repr += f"ignore_index={self.ignore_index}, reduction={self.reduction}, normalized={self.normalized}, "
+        repr += f"reduced_threshold={self.reduced_threshold}, activation={self.activation}, "
+        repr += f"softmax_dim={self.softmax_dim})"
+        return repr
+
     def forward(self, inputs: Tensor, targets: Tensor) -> Tensor:
         """
         Compute focal loss for binary classification problem.

pytorch_toolbelt/losses/functional.py

@@ -16,6 +16,7 @@
 ]
 
 
+@torch.cuda.amp.autocast(False)
 def focal_loss_with_logits(
     output: torch.Tensor,
     target: torch.Tensor,
@@ -53,7 +54,8 @@ def focal_loss_with_logits(
     References:
         https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/loss/losses.py
     """
-    target = target.type_as(output)
+    output = output.float()
+    target = target.float()
 
     if activation == "sigmoid":
         p = torch.sigmoid(output)
@@ -88,9 +90,9 @@ def focal_loss_with_logits(
     if reduction == "mean":
         loss = loss.mean()
     if reduction == "sum":
-        loss = loss.sum(dtype=torch.float32)
+        loss = loss.sum()
     if reduction == "batchwise_mean":
-        loss = loss.sum(dim=0, dtype=torch.float32)
+        loss = loss.sum(dim=0)
 
     return loss