How to force model to focus on "rib region" of Chest X-Ray?

82 Views Asked by At

I have built a model to classify if a Chest X-ray have Rib-Fracture (label with 1) or not (label with 0).

I have used Integrated Gradients - An attribution method scores the input data (the input's pixels in this case) based on the predictions the model makes.

However, based on the attribution maps, the attribution scores for the rib regions is too low, i.e the model does not focus on the rib regions - which is the region model need to "look at" to make precision decisions. An example is shown as below.

A Chest X-ray (left) and its corresponding attribution score map (right)

My question is, is there any ways to force the model focus on the rib regions?.

The model I have used is torchvision.models.densenet161 as Feature Extractor and a classifier. The model's architecture as shown below.

class DenseNet161(nn.Module):
    """Model modified.
    The architecture of our model is the same as standard DenseNet161
    except the classifier layer which has an additional sigmoid function.
    """
    def __init__(self, out_size):
        super(DenseNet161, self).__init__()
        # self.densenet161 = torchvision.models.densenet161(pretrained = False) # `pretrained=False` is deprecated
        self.weights = torchvision.models.DenseNet161_Weights.DEFAULT
        self.densenet161 = torchvision.models.densenet161(weights=self.weights)
        num_ftrs = self.densenet161.classifier.in_features
        self.densenet161.classifier = nn.Sequential(
            nn.Linear(num_ftrs, out_size),
#             nn.Sigmoid()
        )

    def forward(self, x, mode='inference'):
        if mode == 'train':
            x = self.densenet161(x)
        else:
            x = self.densenet161(x)
            x = nn.Sigmoid()(x)
        return x

I have used CBAM between dense_layer inside dense_block as well as before transition_block to help the model focus on more useful features (better with ribs). The performance is improved but not to much with the attribution score map is still not giving high score to rib regions as before.

Also my planning is:

  1. Use the segmentation model to segment the ribs in image, then feed them to model.
  2. Use a Rib Fracture Detection Model.

I have read a lot papers of these 2 methods, but the authors all used the private dataset from hospitals. So I think the only way is somehow guild my model to focus on the regions which contain ribs.

0

There are 0 best solutions below