diff --git a/src/deepness/common/processing_parameters/classify_chip_parameters.py b/src/deepness/common/processing_parameters/classify_chip_parameters.py new file mode 100644 index 00000000..1bb34786 --- /dev/null +++ b/src/deepness/common/processing_parameters/classify_chip_parameters.py @@ -0,0 +1,17 @@ +from dataclasses import dataclass +from typing import Optional + +from deepness.common.processing_parameters.map_processing_parameters import MapProcessingParameters + + +@dataclass +class ClassifyChipParameters(MapProcessingParameters): + """ + Parameters for Classifying Chips obtained from UI. + """ + + raster_id: Optional[str] # id for map layer + vector_id: Optional[str] # id for bounding box layer + config: dict # model configuration + # config keys: 'model_name', 'weights_ckpt', 'class_names', 'normalization_mean', + # 'normalization_std', 'model_arch', 'n_channels', 'image_size' \ No newline at end of file diff --git a/src/deepness/common/processing_parameters/detection_parameters.py b/src/deepness/common/processing_parameters/detection_parameters.py index c96232c6..972be71f 100644 --- a/src/deepness/common/processing_parameters/detection_parameters.py +++ b/src/deepness/common/processing_parameters/detection_parameters.py @@ -2,8 +2,6 @@ from dataclasses import dataclass from deepness.common.processing_parameters.map_processing_parameters import MapProcessingParameters -from deepness.processing.models.model_base import ModelBase - @dataclass class DetectorTypeParameters: @@ -24,6 +22,7 @@ class DetectorType(enum.Enum): YOLO_ULTRALYTICS = 'YOLO_Ultralytics' YOLO_ULTRALYTICS_SEGMENTATION = 'YOLO_Ultralytics_segmentation' YOLO_ULTRALYTICS_OBB = 'YOLO_Ultralytics_obb' + RT_DETR = 'RT_DETR' def get_parameters(self): if self == DetectorType.YOLO_v5_v7_DEFAULT: @@ -42,6 +41,10 @@ def get_parameters(self): has_inverted_output_shape=True, skipped_objectness_probability=True, ) + elif self == DetectorType.RT_DETR: + return DetectorTypeParameters( + skipped_objectness_probability=True, + ) else: raise ValueError(f'Unknown detector type: {self}') @@ -61,7 +64,7 @@ class DetectionParameters(MapProcessingParameters): """ Parameters for Inference of detection model (including pre/post-processing) obtained from UI. """ - + from deepness.processing.models.model_base import ModelBase model: ModelBase # wrapper of the loaded model confidence: float diff --git a/src/deepness/deepness.py b/src/deepness/deepness.py index 9349af48..7362ff4c 100644 --- a/src/deepness/deepness.py +++ b/src/deepness/deepness.py @@ -8,9 +8,9 @@ import logging import traceback -from qgis.core import Qgis, QgsApplication, QgsProject, QgsVectorLayer +from qgis.core import Qgis, QgsApplication, QgsProject, QgsVectorLayer, QgsField, QgsWkbTypes from qgis.gui import QgisInterface -from qgis.PyQt.QtCore import QCoreApplication, Qt +from qgis.PyQt.QtCore import QCoreApplication, Qt, QVariant from qgis.PyQt.QtGui import QIcon from qgis.PyQt.QtWidgets import QAction, QMessageBox @@ -18,6 +18,8 @@ from deepness.common.lazy_package_loader import LazyPackageLoader from deepness.common.processing_parameters.map_processing_parameters import MapProcessingParameters, ProcessedAreaType from deepness.common.processing_parameters.training_data_export_parameters import TrainingDataExportParameters +from deepness.common.processing_parameters.classify_chip_parameters import ClassifyChipParameters + from deepness.deepness_dockwidget import DeepnessDockWidget from deepness.dialogs.resizable_message_box import ResizableMessageBox from deepness.images.get_image_path import get_icon_path @@ -25,6 +27,8 @@ MapProcessingResultFailed, MapProcessingResultSuccess) from deepness.processing.map_processor.map_processor_training_data_export import MapProcessorTrainingDataExport +from deepness.processing.classify_chip import ClassifyChipTask + cv2 = LazyPackageLoader('cv2') @@ -167,6 +171,9 @@ def onClosePlugin(self): # when closing the docked window: # self.dockwidget = None + # QgsProject.instance().layersAdded.disconnect(self.dockwidget._refresh_bbox_combo) + # QgsProject.instance().layersRemoved.disconnect(self.dockwidget._refresh_bbox_combo) + self.pluginIsActive = False def unload(self): @@ -203,6 +210,7 @@ def run(self): self.dockwidget.closingPlugin.connect(self.onClosePlugin) self.dockwidget.run_model_inference_signal.connect(self._run_model_inference) self.dockwidget.run_training_data_export_signal.connect(self._run_training_data_export) + self.dockwidget.run_classify_chip_signal.connect(self._classify_chip) self.iface.addDockWidget(Qt.RightDockWidgetArea, self.dockwidget) self.dockwidget.show() @@ -250,6 +258,56 @@ def _run_training_data_export(self, training_data_export_parameters: TrainingDat QgsApplication.taskManager().addTask(self._map_processor) self._display_processing_started_info() + def _classify_chip(self, classify_chip_parameters: ClassifyChipParameters): + raster_id = classify_chip_parameters.raster_id + vector_id = classify_chip_parameters.vector_id + cfg = classify_chip_parameters.config + task = ClassifyChipTask(raster_id, vector_id, cfg) + + def on_done(ok, task=task): + if not ok or not task.results: + print("Classification task failed or empty results.") + return + + # temporary memory layer + src = task.vec # already reprojected bboxes + geom_name = QgsWkbTypes.displayString(src.wkbType()) + dup = QgsVectorLayer(f"{geom_name}?crs={src.crs().authid()}", task.OUT_LAYER_NAME, "memory") + + # copy original attributes and features + prov = dup.dataProvider() + prov.addAttributes(src.fields()) + dup.updateFields() + prov.addFeatures(list(src.getFeatures())) + + # add custom attribute fields + dup.startEditing() + dup.addAttribute(QgsField(task.ATTR_LABEL_FIELD, QVariant.String)) # label + for cname in task.CLASS_NAMES: # class probabilities + dup.addAttribute(QgsField(f"{task.ATTR_PROB_PREFIX}{cname}", QVariant.Double)) + dup.updateFields() + + # fill attribute fields by feature id (fid) + for f in dup.getFeatures(): + fid = f.id() + if fid in task.results: + label, probs = task.results[fid] + f[task.ATTR_LABEL_FIELD] = label + for j, cname in enumerate(task.CLASS_NAMES): + f[f"{task.ATTR_PROB_PREFIX}{cname}"] = float(probs[j]) + dup.updateFeature(f) + dup.commitChanges() + + # add temporary layer to QGIS + QgsProject.instance().addMapLayer(dup) + print(f"Added prediction layer: {task.OUT_LAYER_NAME} with fields: " + f"{task.ATTR_LABEL_FIELD}, {[task.ATTR_PROB_PREFIX + c for c in task.CLASS_NAMES]}") + + task.taskCompleted.connect(lambda: on_done(True)) + task.taskTerminated.connect(lambda: on_done(False)) + QgsApplication.taskManager().addTask(task) + self._display_processing_started_info() + def _run_model_inference(self, params: MapProcessingParameters): from deepness.processing.models.model_types import ModelDefinition # import here to avoid pulling external dependencies to early diff --git a/src/deepness/deepness_dockwidget.py b/src/deepness/deepness_dockwidget.py index 24d1409e..bb0f63f8 100644 --- a/src/deepness/deepness_dockwidget.py +++ b/src/deepness/deepness_dockwidget.py @@ -6,10 +6,12 @@ import os from typing import Optional +from qgis.gui import QgsCollapsibleGroupBox from qgis.core import Qgis, QgsMapLayerProxyModel, QgsProject from qgis.PyQt import QtWidgets, uic from qgis.PyQt.QtCore import pyqtSignal from qgis.PyQt.QtWidgets import QComboBox, QFileDialog, QMessageBox +from qgis.PyQt.QtWidgets import QFormLayout, QWidget, QLineEdit, QPushButton, QSpinBox from deepness.common.config_entry_key import ConfigEntryKey from deepness.common.defines import IS_DEBUG, PLUGIN_NAME @@ -23,6 +25,7 @@ from deepness.common.processing_parameters.segmentation_parameters import SegmentationParameters from deepness.common.processing_parameters.superresolution_parameters import SuperresolutionParameters from deepness.common.processing_parameters.training_data_export_parameters import TrainingDataExportParameters +from deepness.common.processing_parameters.classify_chip_parameters import ClassifyChipParameters from deepness.processing.models.model_base import ModelBase from deepness.widgets.input_channels_mapping.input_channels_mapping_widget import InputChannelsMappingWidget from deepness.widgets.training_data_export_widget.training_data_export_widget import TrainingDataExportWidget @@ -43,6 +46,7 @@ class DeepnessDockWidget(QtWidgets.QDockWidget, FORM_CLASS): closingPlugin = pyqtSignal() run_model_inference_signal = pyqtSignal(MapProcessingParameters) # run Segmentation or Detection run_training_data_export_signal = pyqtSignal(TrainingDataExportParameters) + run_classify_chip_signal = pyqtSignal(ClassifyChipParameters) def __init__(self, iface, parent=None): super(DeepnessDockWidget, self).__init__(parent) @@ -142,6 +146,57 @@ def _save_ui_to_config(self): def _rlayer_updated(self): self._input_channels_mapping_widget.set_rlayer(self._get_input_layer()) + + ####################### Helper functions for classification begins ####################### + def _mk_layer_combo(self, filter_model): + combo = self._new_layer_combo(filter_model) + return combo + + def _new_layer_combo(self, filter_model): + c = self.mMapLayerComboBox_areaMaskLayer.__class__(self) # duplicate widget + c.setFilters(filter_model) # filter for vector layers (bounding box layers) + return c + + # def _refresh_bbox_combo(self, *args): + # # QgsMapLayerComboBox refreshes automatically; nothing to do here + # pass + + def _pick_ckpt(self, lineedit): + fp, _ = QFileDialog.getOpenFileName(self, "Pick weights checkpoint", "", "PyTorch (*.pt *.pth);;All (*)") + if fp: + lineedit.setText(fp) + + def _parse_strings(self, s): + return [x.strip() for x in s.split(",") if x.strip()] + + def _parse_floats(self, s): + return [float(x.strip()) for x in s.split(",") if x.strip()] + + def _emit_classify_chip(self): + raster = self.mMapLayerComboBox_inputLayer.currentLayer() # pre-selected input data + bbox = self.cmb_bbox.currentLayer() # list of bounding box layers + if raster is None or bbox is None: + QMessageBox.warning(self, "Deepness", "Pick raster and vector (bounding box) layer.") + return + cfg = { + 'model_name': None, + 'weights_ckpt': self.le_ckpt.text(), + 'class_names': self._parse_strings(self.le_classes.text()), + 'normalization_mean': self._parse_floats(self.le_mean.text()), + 'normalization_std': self._parse_floats(self.le_std.text()), + 'model_arch': self.le_arch.text(), + 'n_channels': int(self.sb_ch.value()), + 'image_size': int(self.sb_size.value()), + } + map_processing_parameters = self._get_map_processing_parameters() + params = ClassifyChipParameters( + **map_processing_parameters.__dict__, + raster_id= raster.id(), + vector_id=bbox.id(), + config=cfg, + ) + self.run_classify_chip_signal.emit(params) + ####################### Helper functions for classification ends ####################### def _setup_misc_ui(self): """ Setup some misceleounous ui forms @@ -171,6 +226,60 @@ def _setup_misc_ui(self): self.comboBox_detectorType.addItem(detector_type) self._detector_type_changed() + ####################### GUI for classification begins ####################### + # set up panel + self.group_chip = QgsCollapsibleGroupBox("Chip classification on input data selected above", self) + self.group_chip.setCollapsed(False) + form = QFormLayout(self.group_chip) + + # select vector (bounding box) layer + self.cmb_bbox = self._mk_layer_combo(filter_model=QgsMapLayerProxyModel.VectorLayer) + form.addRow("Bounding boxes:", self.cmb_bbox) + + # select path to model checkpoint + self.le_ckpt = QLineEdit('models/d4-cls-exp1_6ch-3cat-48px-newval_best.pt') + btn_ckpt = QPushButton("Browse…") + btn_ckpt.clicked.connect(lambda: self._pick_ckpt(self.le_ckpt)) + row_ckpt = QtWidgets.QHBoxLayout() # arrange widgets left to right + row_ckpt.addWidget(self.le_ckpt) + row_ckpt.addWidget(btn_ckpt) + wrap_ckpt = QWidget() # wrap the row of widgets into one + wrap_ckpt.setLayout(row_ckpt) + form.addRow("Model file path (.pt):", wrap_ckpt) + + # select model configuration (with default values) + self.le_mname = QLineEdit('Tree Classifier') + self.le_classes = QLineEdit('deadtree,topdownthreat,tree') + self.le_mean = QLineEdit('0.358034,0.492683,0.479417,0.502045,0.499618,0.536634') + self.le_std = QLineEdit('0.213879,0.261817,0.220773,0.279437,0.281337,0.199962') + self.le_arch = QLineEdit('tf_efficientnet_b4.ns_jft_in1k') + self.sb_ch = QSpinBox(); self.sb_ch.setRange(1, 12); self.sb_ch.setValue(6) + self.sb_size = QSpinBox(); self.sb_size.setRange(16, 1024); self.sb_size.setValue(48) + + form.addRow("Model Name:", self.le_mname) + form.addRow("Class Names:", self.le_classes) + form.addRow("Normalization Mean:", self.le_mean) + form.addRow("Normalization Std:", self.le_std) + form.addRow("Model architecture:", self.le_arch) + form.addRow("Number of Channels:", self.sb_ch) + form.addRow("Image Size:", self.sb_size) + + # button to run classification + btn_run_chip = QPushButton("Classify Chips") + btn_run_chip.clicked.connect(self._emit_classify_chip) + form.addRow("", btn_run_chip) + + # insert the entire block just below "Training data export" block + export_block = self.verticalLayout_trainingDataExport.parentWidget() + layout = self.verticalLayout_3 + layout.insertWidget(layout.indexOf(export_block)+1, self.group_chip) + + # # keep comboboxes synced with project layers + # QgsProject.instance().layersAdded.connect(self._refresh_bbox_combo) + # QgsProject.instance().layersRemoved.connect(self._refresh_bbox_combo) + # self._refresh_bbox_combo() + ####################### GUI for classification ends ####################### + self._rlayer_updated() # to force refresh the dependant ui elements def _set_processed_area_mask_options(self): diff --git a/src/deepness/metadata.txt b/src/deepness/metadata.txt index 32b8bd13..d150859a 100644 --- a/src/deepness/metadata.txt +++ b/src/deepness/metadata.txt @@ -6,7 +6,7 @@ name=Deepness: Deep Neural Remote Sensing qgisMinimumVersion=3.22 description=Inference of deep neural network models (ONNX) for segmentation, detection and regression -version=0.7.0 +version=0.6.5 author=PUT Vision email=przemyslaw.aszkowski@gmail.com diff --git a/src/deepness/processing/classify_chip.py b/src/deepness/processing/classify_chip.py new file mode 100644 index 00000000..bb42b5f3 --- /dev/null +++ b/src/deepness/processing/classify_chip.py @@ -0,0 +1,126 @@ +# import packages +from qgis.core import QgsTask, QgsApplication, QgsProject, QgsVectorLayer, QgsField, QgsWkbTypes +from qgis.PyQt.QtCore import QVariant +import torch +import numpy as np + +# helper function to initialize model +class ImgClassifier(torch.nn.Module): + def __init__(self, model_arch, n_class, n_channels=3, pretrained=False): + super().__init__() + import timm + self.model = timm.create_model(model_arch, in_chans=n_channels, pretrained=pretrained) + n_features = self.model.classifier.in_features + self.model.classifier = torch.nn.Linear(n_features, n_class) + def forward(self, x): + return self.model(x) + +# QGIS task to extract and classify chips +class ClassifyChipTask(QgsTask): + def __init__(self, raster_layer_id, vector_layer_id, model_config): + super().__init__("Extract & Classify Chips", QgsTask.CanCancel) + print(f'raster_layer_id: {raster_layer_id}') + print(f'vector_layer_id: {vector_layer_id}') + print(f'model_config: {model_config}') + + self.ras = QgsProject.instance().mapLayer(raster_layer_id) + self.vec = QgsProject.instance().mapLayer(vector_layer_id) + + self.MODEL_NAME = model_config['model_name'] + self.CKPT = model_config['weights_ckpt'] + self.CLASS_NAMES = model_config['class_names'] + self.MEAN = model_config['normalization_mean'] + self.STD = model_config['normalization_std'] + self.MODEL_ARCH = model_config['model_arch'] + self.N_CHANNELS = model_config['n_channels'] + self.NUM_CLASSES = len(self.CLASS_NAMES) + self.IMG_SIZE = model_config['image_size'] + + if not self.MODEL_NAME: + self.MODEL_NAME = self.CKPT.replace('\\','/').split('/')[-1].split('.')[0] + + self.ATTR_LABEL_FIELD = "label" # string label + self.ATTR_PROB_PREFIX = "p_" # p_deadtree, p_topdownthreat, p_tree + self.OUT_LAYER_NAME = f"{self.MODEL_NAME}_classification" + self.results = {} # key: feature id; val: (label, probs) + + def _prep_model(self): + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + model = ImgClassifier(self.MODEL_ARCH, self.NUM_CLASSES, n_channels=self.N_CHANNELS, pretrained=False) + model.load_state_dict(torch.load(self.CKPT, map_location='cpu')) + model.to(device).eval() + return model, device + + def _read_chip_np(self, src_ds, xoff, yoff, xsize, ysize, device): + from albumentations import Compose, Resize, Normalize + from albumentations.pytorch import ToTensorV2 + arr = src_ds.ReadAsArray(xoff, yoff, xsize, ysize) # shape: (C, H, W) + arr = np.transpose(arr, (1, 2, 0)) # shape: (H, W, C) + max_pixel_value = 255.0 if arr.max() > 1.5 else 1.0 + tfm = Compose([ + Resize(self.IMG_SIZE, self.IMG_SIZE), + Normalize(mean=self.MEAN, std=self.STD, max_pixel_value=max_pixel_value), + ToTensorV2() + ]) + x = tfm(image=arr)['image'] + return x.unsqueeze(0).float().to(device) + + def run(self): + import time + start_time = time.time() + # 1) reproject vector to match raster CRS + if self.vec.crs() != self.ras.crs(): + import processing + self.vec = processing.run("native:reprojectlayer", { + "INPUT": self.vec, "TARGET_CRS": self.ras.crs(), "OUTPUT": "memory:" + })["OUTPUT"] + + # 2) load Raster grid extent and source + extent = self.ras.extent() + px_w = extent.width() / self.ras.width() + px_h = extent.height() / self.ras.height() + origin_x = extent.xMinimum() + origin_y = extent.yMaximum() + + from osgeo import gdal + src_ds = gdal.Open(self.ras.source(), gdal.GA_ReadOnly) + + # 3) extract and classify chips + model, device = self._prep_model() + feats = list(self.vec.getFeatures()) + total = len(feats) + + with torch.no_grad(): + for i, f in enumerate(feats): + if self.isCanceled(): + break + + # extract bounding box + e = f.geometry().boundingBox() + xoff = int((e.xMinimum() - origin_x) / px_w) + yoff = int((origin_y - e.yMaximum()) / px_h) + xsize = int(e.width() / px_w) + ysize = int(e.height() / px_h) + if xsize <= 0 or ysize <= 0: + continue + + # read chip + chip_np = self._read_chip_np(src_ds, xoff, yoff, xsize, ysize, device) + + # classify chip + with torch.no_grad(): + logits = model(chip_np) + probs = torch.softmax(logits, dim=1).squeeze(0).cpu().numpy() + label = self.CLASS_NAMES[int(probs.argmax())] + + # save results + self.results[f.id()] = (label, probs.tolist()) + self.setProgress(100.0 * (i + 1) / max(1, total)) + + src_ds = None + end_time = time.time() + print(f'Processing completed; time taken {round((end_time-start_time)/60,2)} minutes') + return True + + def finished(self, ok): + pass \ No newline at end of file diff --git a/src/deepness/processing/map_processor/map_processor_detection.py b/src/deepness/processing/map_processor/map_processor_detection.py index 8b5658b4..83c60b9b 100644 --- a/src/deepness/processing/map_processor/map_processor_detection.py +++ b/src/deepness/processing/map_processor/map_processor_detection.py @@ -4,8 +4,6 @@ import cv2 import numpy as np from qgis.core import QgsFeature, QgsGeometry, QgsProject, QgsVectorLayer -from qgis.PyQt.QtCore import QVariant -from qgis.core import QgsFields, QgsField from deepness.common.processing_parameters.detection_parameters import DetectionParameters from deepness.processing import processing_utils @@ -45,11 +43,13 @@ def get_all_detections(self) -> List[Detection]: def _run(self) -> MapProcessingResult: all_bounding_boxes = [] # type: List[Detection] + i=0 for tile_img_batched, tile_params_batched in self.tiles_generator_batched(): if self.isCanceled(): return MapProcessingResultCanceled() - bounding_boxes_in_tile_batched = self._process_tile(tile_img_batched, tile_params_batched) + bounding_boxes_in_tile_batched = self._process_tile(tile_img_batched, tile_params_batched, i) + i+=1 all_bounding_boxes += [d for det in bounding_boxes_in_tile_batched for d in det] with_rot = self.detection_parameters.detector_type == DetectorType.YOLO_ULTRALYTICS_OBB @@ -61,9 +61,9 @@ def _run(self) -> MapProcessingResult: all_bounding_boxes_restricted = [] gui_delegate = self._create_vlayer_for_output_bounding_boxes(all_bounding_boxes_restricted) - result_message = self._create_result_message(all_bounding_boxes_restricted) self._all_detections = all_bounding_boxes_restricted + return MapProcessingResultSuccess( message=result_message, gui_delegate=gui_delegate, @@ -130,16 +130,9 @@ def _create_vlayer_for_output_bounding_boxes(self, bounding_boxes: List[Detectio for channel_id in channels: filtered_bounding_boxes = [det for det in bounding_boxes if det.clss == channel_id] print(f'Detections for class {channel_id}: {len(filtered_bounding_boxes)}') - - vlayer = QgsVectorLayer("multipolygon", self.model.get_channel_name(0, channel_id), "memory") - vlayer.setCrs(self.rlayer.crs()) - prov = vlayer.dataProvider() - prov.addAttributes([QgsField("confidence", QVariant.Double)]) - vlayer.updateFields() features = [] for det in filtered_bounding_boxes: - feature = QgsFeature() if det.mask is None: bbox_corners_pixels = det.bbox.get_4_corners() bbox_corners_crs = processing_utils.transform_points_list_xy_to_target_crs( @@ -147,13 +140,13 @@ def _create_vlayer_for_output_bounding_boxes(self, bounding_boxes: List[Detectio extent=self.extended_extent, rlayer_units_per_pixel=self.rlayer_units_per_pixel, ) - #feature = QgsFeature() #move outside of the if block + feature = QgsFeature() polygon_xy_vec_vec = [ bbox_corners_crs ] geometry = QgsGeometry.fromPolygonXY(polygon_xy_vec_vec) - #feature.setGeometry(geometry) - #features.append(feature) + feature.setGeometry(geometry) + features.append(feature) else: contours, _ = cv2.findContours(det.mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) contours = sorted(contours, key=cv2.contourArea, reverse=True) @@ -176,20 +169,17 @@ def _create_vlayer_for_output_bounding_boxes(self, bounding_boxes: List[Detectio rlayer_units_per_pixel=self.rlayer_units_per_pixel, ) - #feature = QgsFeature() + feature = QgsFeature() polygon_xy_vec_vec = [ mask_corners_crs ] geometry = QgsGeometry.fromPolygonXY(polygon_xy_vec_vec) - #feature.setGeometry(geometry) - #features.append(feature) - feature.setGeometry(geometry) - feature.setAttributes([float(det.conf)]) - features.append(feature) - - #vlayer = QgsVectorLayer("multipolygon", self.model.get_channel_name(0, channel_id), "memory") - #vlayer.setCrs(self.rlayer.crs()) - #prov = vlayer.dataProvider() + feature.setGeometry(geometry) + features.append(feature) + + vlayer = QgsVectorLayer("multipolygon", self.model.get_channel_name(0, channel_id), "memory") + vlayer.setCrs(self.rlayer.crs()) + prov = vlayer.dataProvider() color = vlayer.renderer().symbol().color() OUTPUT_VLAYER_COLOR_TRANSPARENCY = 80 @@ -279,8 +269,8 @@ def convert_bounding_boxes_to_absolute_positions(bounding_boxes_relative: List[D for det in bounding_boxes_relative: det.convert_to_global(offset_x=tile_params.start_pixel_x, offset_y=tile_params.start_pixel_y) - def _process_tile(self, tile_img: np.ndarray, tile_params_batched: List[TileParams]) -> np.ndarray: - bounding_boxes_batched: List[Detection] = self.model.process(tile_img) + def _process_tile(self, tile_img: np.ndarray, tile_params_batched: List[TileParams],i) -> np.ndarray: + bounding_boxes_batched: List[Detection] = self.model.process(tile_img, i) for bounding_boxes, tile_params in zip(bounding_boxes_batched, tile_params_batched): self.convert_bounding_boxes_to_absolute_positions(bounding_boxes, tile_params) diff --git a/src/deepness/processing/models/detector.py b/src/deepness/processing/models/detector.py index 474802d5..27d18090 100644 --- a/src/deepness/processing/models/detector.py +++ b/src/deepness/processing/models/detector.py @@ -174,7 +174,7 @@ def get_number_of_output_channels(self): else: raise NotImplementedError("Model with multiple output layer is not supported! Use only one output layer.") - def postprocessing(self, model_output): + def postprocessing(self, model_output, TILE_SIZE): """Postprocess model output NOTE: Maybe refactor this, as it has many added layers of checks which can be simplified. @@ -201,8 +201,12 @@ def postprocessing(self, model_output): batch_detection = [] outputs_range = len(model_output) - - if self.model_type == DetectorType.YOLO_ULTRALYTICS_SEGMENTATION or self.model_type == DetectorType.YOLO_v9: + + if (self.model_type == DetectorType.YOLO_ULTRALYTICS_SEGMENTATION or + self.model_type == DetectorType.YOLO_v9): + outputs_range = len(model_output[0]) + + if (self.model_type == DetectorType.RT_DETR): outputs_range = len(model_output[0]) for i in range(outputs_range): @@ -222,6 +226,8 @@ def postprocessing(self, model_output): boxes, conf, classes, masks = self._postprocessing_YOLO_ULTRALYTICS_SEGMENTATION(model_output[0][i], model_output[1][i]) elif self.model_type == DetectorType.YOLO_ULTRALYTICS_OBB: boxes, conf, classes, rots = self._postprocessing_YOLO_ULTRALYTICS_OBB(model_output[0][i]) + elif self.model_type == DetectorType.RT_DETR: + boxes, conf, classes = self._postprocessing_RT_DETR(model_output[0][i], TILE_SIZE) else: raise NotImplementedError(f"Model type not implemented! ('{self.model_type}')") @@ -322,6 +328,35 @@ def _postprocessing_YOLO_v9(self, model_output): classes = np.argmax(outputs_nms[:, 4:], axis=1) return boxes, conf, classes + + def _postprocessing_RT_DETR(self, model_output, TILE_SIZE): + # filter outputs by confidence score + outputs_filtered = np.array(list(filter(lambda x: np.max(x[4:]) >= self.confidence, + model_output))) + if len(outputs_filtered.shape) < 2: + return [], [], [] + outputs_x1y1x2y2 = self.xywh2xyxy(outputs_filtered) + boxes = outputs_x1y1x2y2[:,:4] + probabilities = outputs_x1y1x2y2[:, 4:][:,0] + + # scale boudning boxes + boxes = boxes * TILE_SIZE + boxes = np.array(boxes, dtype=int) + + # nms + pick_indxs = self.non_max_suppression_fast( + boxes, + probs=probabilities, + iou_threshold=self.iou_threshold) + + boxes = boxes[pick_indxs] + probabilities = probabilities[pick_indxs] + + boxes = np.array(boxes, dtype=int) + conf = probabilities + classes = np.argmax(np.expand_dims(conf, axis=1), axis=1) + + return boxes, conf, classes def _postprocessing_YOLO_ULTRALYTICS(self, model_output): model_output = np.transpose(model_output, (1, 0)) diff --git a/src/deepness/processing/models/model_base.py b/src/deepness/processing/models/model_base.py index 11dc214f..acfe7043 100644 --- a/src/deepness/processing/models/model_base.py +++ b/src/deepness/processing/models/model_base.py @@ -7,6 +7,7 @@ from deepness.common.lazy_package_loader import LazyPackageLoader from deepness.common.processing_parameters.standardization_parameters import StandardizationParameters +from deepness.common.processing_parameters.detection_parameters import DetectorType ort = LazyPackageLoader('onnxruntime') @@ -47,6 +48,7 @@ def __init__(self, model_file_path: str): self.standardization_parameters: StandardizationParameters = self.get_metadata_standarization_parameters() self.outputs_names = self.get_outputs_channel_names() + self.model_type = '' @classmethod def get_model_type_from_metadata(cls, model_file_path: str) -> Optional[str]: @@ -368,7 +370,7 @@ def get_number_of_channels(self) -> int: """ return self.input_shape[-3] - def process(self, tiles_batched: np.ndarray): + def process(self, tiles_batched: np.ndarray, i=None): """ Process a single tile image Parameters @@ -381,11 +383,16 @@ def process(self, tiles_batched: np.ndarray): np.ndarray Single prediction """ + TILE_SIZE = tiles_batched.shape[1] + # print(f'unprocessed tiles_batched.shape: {tiles_batched.shape}') + # np.save(f"unprocessed_qgis_input{i}", tiles_batched) input_batch = self.preprocessing(tiles_batched) + # print(f'preprocessed tiles_batched.shape: {input_batch.shape}') + # np.save(f"preprocessed_qgis_input{i}", input_batch) model_output = self.sess.run( output_names=None, input_feed={self.input_name: input_batch}) - res = self.postprocessing(model_output) + res = self.postprocessing(model_output, TILE_SIZE) return res def preprocessing(self, tiles_batched: np.ndarray) -> np.ndarray: @@ -407,9 +414,12 @@ def preprocessing(self, tiles_batched: np.ndarray) -> np.ndarray: import deepness.processing.models.preprocessing_utils as preprocessing_utils tiles_batched = preprocessing_utils.limit_channels_number(tiles_batched, limit=self.input_shape[-3]) - tiles_batched = preprocessing_utils.normalize_values_to_01(tiles_batched) - tiles_batched = preprocessing_utils.standardize_values(tiles_batched, params=self.standardization_parameters) - tiles_batched = preprocessing_utils.transpose_nhwc_to_nchw(tiles_batched) + if self.model_type != DetectorType.RT_DETR: + tiles_batched = preprocessing_utils.normalize_values_to_01(tiles_batched) + tiles_batched = preprocessing_utils.standardize_values(tiles_batched, params=self.standardization_parameters) + tiles_batched = preprocessing_utils.transpose_nhwc_to_nchw(tiles_batched) + else: + tiles_batched = preprocessing_utils.normalize_6_channels(tiles_batched) # returns NCHW return tiles_batched diff --git a/src/deepness/processing/models/preprocessing_utils.py b/src/deepness/processing/models/preprocessing_utils.py index eff25fa1..2a709634 100644 --- a/src/deepness/processing/models/preprocessing_utils.py +++ b/src/deepness/processing/models/preprocessing_utils.py @@ -12,6 +12,25 @@ def limit_channels_number(tiles_batched: np.array, limit: int) -> np.array: """ return tiles_batched[:, :, :, :limit] +def normalize_band(band): + """Normalize a single band with contrast stretching, handling bad values.""" + valid_mask = np.isfinite(band) & (band > -1e10) + if not np.any(valid_mask): + return np.zeros_like(band, dtype=np.uint8) + + p2, p98 = np.percentile(band[valid_mask], (2, 98)) + norm = np.clip((band - p2) / (p98 - p2), 0, 1) + return (norm * 255).astype(np.float32) + +def normalize_6_channels(tiles_batched): + # tiles_batched shape: (4,256,256,6) + norm_tiles = [] + for tile_index in range(tiles_batched.shape[0]): + tile = tiles_batched[tile_index] # shape: (256,256,6) + norm_tile = [normalize_band(tile[:,:,i]) for i in range(tile.shape[2])] + norm_tile = np.stack(norm_tile, axis=0) / 255 # shape: (6,256,256) + norm_tiles.append(norm_tile) + return np.stack(norm_tiles, axis=0) # shape: (4,6,256,256) NCHW def normalize_values_to_01(tiles_batched: np.array) -> np.array: """ Normalize the values of the input image to the model to the range [0, 1] @@ -21,7 +40,6 @@ def normalize_values_to_01(tiles_batched: np.array) -> np.array: """ return np.float32(tiles_batched * 1./255.) - def standardize_values(tiles_batched: np.array, params: StandardizationParameters) -> np.array: """ Standardize the input image to the model @@ -29,6 +47,9 @@ def standardize_values(tiles_batched: np.array, params: StandardizationParameter :param params: Parameters for standardization of type STANDARIZE_PARAMS :return: Batch of tiles with standardized values """ + print(f'params.mean: {params.mean}') + print(f'params.std: {params.std}') + return (tiles_batched - params.mean) / params.std