Add parameters to make custom backbone for detr

src/transformers/image_utils.py

@@ -177,8 +177,10 @@ def normalize(self, image, mean, std):
 
             if not isinstance(mean, torch.Tensor):
                 mean = torch.tensor(mean)
+            mean = mean.to(image.device)
             if not isinstance(std, torch.Tensor):
                 std = torch.tensor(std)
+            std = std.to(image.device)
 
         if image.ndim == 3 and image.shape[0] in [1, 3]:
             return (image - mean[:, None, None]) / std[:, None, None]

src/transformers/models/detr/configuration_detr.py

@@ -155,6 +155,10 @@ def __init__(
         bbox_loss_coefficient=5,
         giou_loss_coefficient=2,
         eos_coefficient=0.1,
+        num_channels=3,
+        use_pretrained_backbone=True,
+        freeze_layers=True,
+        fix_batch_norm=True,
         **kwargs
     ):
         self.num_queries = num_queries
@@ -190,6 +194,10 @@ def __init__(
         self.bbox_loss_coefficient = bbox_loss_coefficient
         self.giou_loss_coefficient = giou_loss_coefficient
         self.eos_coefficient = eos_coefficient
+        self.num_channels = num_channels
+        self.use_pretrained_backbone = use_pretrained_backbone
+        self.freeze_layers = freeze_layers
+        self.fix_batch_norm = fix_batch_norm
         super().__init__(is_encoder_decoder=is_encoder_decoder, **kwargs)
 
     @property

src/transformers/models/detr/feature_extraction_detr.py

@@ -581,17 +581,27 @@ def __call__(
         if pad_and_return_pixel_mask:
             # pad images up to largest image in batch and create pixel_mask
             max_size = self._max_by_axis([list(image.shape) for image in images])
-            c, h, w = max_size
+            if len(max_size) != 2:
+                c, h, w = max_size
+            else:
+                h, w = max_size
+                c = 1
             padded_images = []
             pixel_mask = []
             for image in images:
                 # create padded image
                 padded_image = np.zeros((c, h, w), dtype=np.float32)
-                padded_image[: image.shape[0], : image.shape[1], : image.shape[2]] = np.copy(image)
+                if c != 1:
+                    padded_image[: image.shape[0], : image.shape[1], : image.shape[2]] = np.copy(image)
+                else:
+                    padded_image[0, : image.shape[0], : image.shape[1]] = np.copy(image)
                 padded_images.append(padded_image)
                 # create pixel mask
                 mask = np.zeros((h, w), dtype=np.int64)
-                mask[: image.shape[1], : image.shape[2]] = True
+                if c == 1:
+                    mask[: image.shape[0], : image.shape[1]] = True
+                else:
+                    mask[: image.shape[1], : image.shape[2]] = True
                 pixel_mask.append(mask)
             images = padded_images
 

src/transformers/models/detr/modeling_detr.py

@@ -315,23 +315,26 @@ class DetrTimmConvEncoder(nn.Module):
 
     """
 
-    def __init__(self, name: str, dilation: bool):
+    def __init__(
+        self, name: str, dilation: bool, num_channels: int, pretrained=True, freeze_layers=True, fix_batch_norm=True
+    ):
         super().__init__()
 
-        kwargs = {}
+        kwargs = {"in_chans": num_channels}
         if dilation:
             kwargs["output_stride"] = 16
 
         requires_backends(self, ["timm"])
 
-        backbone = create_model(name, pretrained=True, features_only=True, out_indices=(1, 2, 3, 4), **kwargs)
+        backbone = create_model(name, pretrained=pretrained, features_only=True, out_indices=(1, 2, 3, 4), **kwargs)
         # replace batch norm by frozen batch norm
-        with torch.no_grad():
-            replace_batch_norm(backbone)
+        if fix_batch_norm:
+            with torch.no_grad():
+                replace_batch_norm(backbone)
         self.model = backbone
         self.intermediate_channel_sizes = self.model.feature_info.channels()
 
-        if "resnet" in name:
+        if "resnet" in name and freeze_layers:
             for name, parameter in self.model.named_parameters():
                 if "layer2" not in name and "layer3" not in name and "layer4" not in name:
                     parameter.requires_grad_(False)
@@ -1159,7 +1162,14 @@ def __init__(self, config: DetrConfig):
         super().__init__(config)
 
         # Create backbone + positional encoding
-        backbone = DetrTimmConvEncoder(config.backbone, config.dilation)
+        backbone = DetrTimmConvEncoder(
+            config.backbone,
+            config.dilation,
+            config.num_channels,
+            pretrained=config.use_pretrained_backbone,
+            freeze_layers=config.freeze_layers,
+            fix_batch_norm=config.fix_batch_norm,
+        )
         position_embeddings = build_position_encoding(config)
         self.backbone = DetrConvModel(backbone, position_embeddings)