Update models/grasp_mods.py

models/grasp_mods.py

@@ -68,25 +68,65 @@ def modify_grasp_loss_forward(self):
 
         return losses
 
-    def modified_loss_boxes(outputs, targets, indices, num_boxes):
+    def modified_loss_boxes(outputs, targets, indices, num_boxes, ignore_wh=False):
 
         if "pred_boxes" not in outputs:
             raise KeyError("No predicted boxes found in outputs")
         idx = self._get_source_permutation_idx(indices)
         source_boxes = outputs["pred_boxes"][idx]
         target_boxes = torch.cat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0)
-
-        loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none")
+        if not ignore_wh:
+            loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none")
+        else:
+            source_xytheta = source_boxes[:, [0, 1, 4]]
+            target_xytheta = target_boxes[:, [0, 1, 4]]
+            loss_bbox = nn.functional.l1_loss(source_xytheta, target_xytheta, reduction="none") * 5 / 3
 
         losses = {}
         losses["loss_bbox"] = loss_bbox.sum() / num_boxes
-
-        loss_giou = 1 - torch.diag(
-            generalized_box_iou(center_to_corners_format(source_boxes[:, :4]), center_to_corners_format(target_boxes[:, :4]))
-        )
+        if not ignore_wh:
+            loss_giou = 1 - torch.diag(
+                generalized_box_iou(center_to_corners_format(source_boxes[:, :4]), center_to_corners_format(target_boxes[:, :4]))
+            )
+        else:
+            source_boxes[:, -2:] = target_boxes[:, -2:].clone()
+            source_corners = center_to_corners_format(source_boxes[:, :4])
+            target_corners = center_to_corners_format(target_boxes[:, :4])
+            loss_giou = 1 - torch.diag(generalized_box_iou(source_corners, target_corners))
         losses["loss_giou"] = loss_giou.sum() / num_boxes
         return losses
-    return modified_loss_labels, modified_loss_boxes
+    def modified_forward(outputs, targets, ignore_wh=False):
+        """
+        This performs the loss computation.
+
+        Args:
+             outputs (`dict`, *optional*):
+                Dictionary of tensors, see the output specification of the model for the format.
+             targets (`List[dict]`, *optional*):
+                List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the
+                losses applied, see each loss' doc.
+        """
+        outputs_without_aux = {k: v for k, v in outputs.items() if k != "auxiliary_outputs"}
+
+        # Retrieve the matching between the outputs of the last layer and the targets
+        indices = self.matcher(outputs_without_aux, targets)
+
+        # Compute the average number of target boxes across all nodes, for normalization purposes
+        num_boxes = sum(len(t["class_labels"]) for t in targets)
+        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
+        # (Niels): comment out function below, distributed training to be added
+        # if is_dist_avail_and_initialized():
+        #     torch.distributed.all_reduce(num_boxes)
+        # (Niels) in original implementation, num_boxes is divided by get_world_size()
+        num_boxes = torch.clamp(num_boxes, min=1).item()
+
+        # Compute all the requested losses
+        losses = {}
+        losses.update(self.loss_labels(outputs, targets, indices, num_boxes))
+        losses.update(self.loss_boxes(outputs, targets, indices, num_boxes, ignore_wh))
+
+        return losses
+    return modified_loss_labels, modified_loss_boxes, modified_forward
 
 def modify_forward(self):
     """
@@ -127,7 +167,7 @@ def modify_forward(self):
     ):
         input_images = torch.stack([x["image"] for x in batched_input], dim=0)
         image_embeddings = self.image_encoder(input_images)
-
+        batch_size = len(batched_input)
         outputs = []
         srcs = []
         for image_record, curr_embedding in zip(batched_input, image_embeddings):
@@ -162,13 +202,17 @@ def modify_forward(self):
         grasp_query_pe = self.grasp_query_position_embeddings(torch.arange(20).to(self.device))
         # repeat to batchsize
         grasp_query_pe = grasp_query_pe.repeat(len(batched_input), 1, 1)
+        pixel_masks = torch.cat([batched_input[i]['pixel_mask'] for i in range(len(batched_input))], dim=0)
+        downsampled_pixel_masks = nn.functional.interpolate(pixel_masks.unsqueeze(1).float(), size=(64, 64), mode='nearest').squeeze(1).bool()
+        downsampled_pixel_masks = downsampled_pixel_masks.view(batch_size, 64*64).contiguous()
+        grasp_encoder_hidden_states = grasp_encoder_hidden_states.view(batch_size, 64*64, 256).contiguous()
         grasp_decoder_outputs = self.grasp_decoder(
             inputs_embeds=torch.zeros_like(grasp_query_pe),
             attention_mask=None,
             position_embeddings=torch.zeros_like(grasp_encoder_hidden_states),
             query_position_embeddings=grasp_query_pe,
             encoder_hidden_states=grasp_encoder_hidden_states,
-            encoder_attention_mask=None,
+            encoder_attention_mask=downsampled_pixel_masks,
             output_attentions=False,
             output_hidden_states=False,
             return_dict=True,
@@ -198,14 +242,14 @@ def modify_forward(self):
                 eos_coef=config.eos_coefficient,
                 losses=losses,
             )
-            criterion.loss_labels, criterion.loss_boxes = modify_grasp_loss_forward(criterion)
+            criterion.loss_labels, criterion.loss_boxes, criterion.forward = modify_grasp_loss_forward(criterion)
             criterion.to(self.device)
             # Third: compute the losses, based on outputs and labels
             outputs_loss = {}
             outputs_loss["logits"] = grasp_logits
             outputs_loss["pred_boxes"] = pred_grasps
 
-            grasp_loss_dict = criterion(outputs_loss, grasp_labels)
+            grasp_loss_dict = criterion(outputs_loss, grasp_labels, ignore_wh=batched_input[0].get("ignore_wh", False))
             # Fourth: compute total loss, as a weighted sum of the various losses
             weight_dict = {"loss_ce": 1, "loss_bbox": config.bbox_loss_coefficient}
             weight_dict["loss_giou"] = config.giou_loss_coefficient
@@ -282,6 +326,8 @@ def add_inference_method(self):
 
         n_queries = iou_predictions.size(0)
 
+        batch_size = n_queries
+
         # forward grasp decoder here
         # 1. Get encoder hidden states
         grasp_encoder_hidden_states = self.grasp_img_pos_embed(src.permute(0, 2, 3, 1))
@@ -289,13 +335,18 @@ def add_inference_method(self):
         grasp_query_pe = self.grasp_query_position_embeddings(torch.arange(20).to(self.device))
         # repeat to batchsize
         grasp_query_pe = grasp_query_pe.repeat(n_queries, 1, 1)
+        pixel_masks = torch.cat([batched_input[i]['pixel_mask'] for i in range(len(batched_input))], dim=0)
+        downsampled_pixel_masks = nn.functional.interpolate(pixel_masks.unsqueeze(1).float(), size=(64, 64),
+                                                            mode='nearest').squeeze(1).bool()
+        downsampled_pixel_masks = downsampled_pixel_masks.view(batch_size, 64 * 64).contiguous()
+        grasp_encoder_hidden_states = grasp_encoder_hidden_states.view(batch_size, 64 * 64, 256).contiguous()
         grasp_decoder_outputs = self.grasp_decoder(
             inputs_embeds=torch.zeros_like(grasp_query_pe),
             attention_mask=None,
             position_embeddings=torch.zeros_like(grasp_encoder_hidden_states),
             query_position_embeddings=grasp_query_pe,
             encoder_hidden_states=grasp_encoder_hidden_states,
-            encoder_attention_mask=None,
+            encoder_attention_mask=downsampled_pixel_masks,
             output_attentions=False,
             output_hidden_states=False,
             return_dict=True,