how can we train model for lines with multi-class labels?

[ballista1704]

If we want to train a model, and we hope this model can not only detect lines in the image, but also give class labels for each line.
How can we modify the code and prepare the dataset?

I think the key might be the num_classes parameter in config files, and the postprocess code in models/linea/linea.py?
Maybe the test(), evaluate() functions will also need to be modified?

class PostProcess(nn.Module):
    """ This module converts the model's output into the format expected by the coco api"""
    def __init__(self) -> None:
        super().__init__()
        self.deploy_mode = False

    @torch.no_grad()
    def forward(self, outputs, target_sizes):
        """ Perform the computation
        Parameters:
            outputs: raw outputs of the model
            target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
                          For evaluation, this must be the original image size (before any data augmentation)
                          For visualization, this should be the image size after data augment, but before padding
        """
        out_logits, out_line = outputs['pred_logits'], outputs['pred_lines']

        """
        Here might be the place to be modified? Change it into something that can handle multi-class scores?
        """
        scores = out_logits[..., 0].sigmoid()

        # convert to [x0, y0, x1, y1] format
        lines = out_line * target_sizes.repeat(1, 2).unsqueeze(1)

        if self.deploy_mode:
            return lines, scores

        results = [{'lines': l, 'scores': s} for s, l in zip(scores, lines)]

        return results

[SebastianJanampa]

Hi @ballista1704 ,
First, i want to apologize for the late response.
You're right. Those are the parameters you have to modify. For the postprocessor, it should be something like this

            scores = F.sigmoid(logits)
            nq = scores.size(1) # number of queries
            scores, index = torch.topk(scores.flatten(1), nq, dim=-1)
            # labels = index % self.num_classes
            labels = mod(index, self.num_classes)
            index = index // self.num_classes