[data.segment] Shape mask based on image as targets can change on inherited classes

src/mednet/data/segment/cxr8.py

@@ -125,9 +125,12 @@ class RawDataLoader(BaseDataLoader):
 
         target = self.target(sample)
 
+        # use image as a base since target() can be overriden by child class
+        mask = torch.ones((1, image.shape[-2], image.shape[-1]), dtype=torch.float32)
+
         image = tv_tensors.Image(image)
         target = tv_tensors.Mask(target)
-        mask = tv_tensors.Mask(torch.ones_like(target))
+        mask = tv_tensors.Mask(mask)
 
         return dict(image=image, target=target, mask=mask, name=sample[0])
 

src/mednet/data/segment/jsrt.py

@@ -141,9 +141,12 @@ class RawDataLoader(BaseDataLoader):
 
         target = self.target(sample)
 
+        # use image as a base since target() can be overriden by child class
+        mask = torch.ones((1, image.shape[-2], image.shape[-1]), dtype=torch.float32)
+
         image = tv_tensors.Image(image)
         target = tv_tensors.Mask(target)
-        mask = tv_tensors.Mask(torch.ones_like(target))
+        mask = tv_tensors.Mask(mask)
 
         return dict(image=image, target=target, mask=mask, name=sample[0])
 

src/mednet/data/segment/montgomery.py

@@ -96,7 +96,8 @@ class RawDataLoader(BaseDataLoader):
 
         target = self.target(sample)
 
-        mask = torch.ones_like(target)
+        # use image as a base since target() can be overriden by child class
+        mask = torch.ones((1, image.shape[-2], image.shape[-1]), dtype=torch.float32)
 
         image = tv_tensors.Image(image)
         target = tv_tensors.Mask(target)

src/mednet/data/segment/shenzhen.py

@@ -113,7 +113,8 @@ class RawDataLoader(BaseDataLoader):
 
         target = self.target(sample)
 
-        mask = torch.ones_like(target)
+        # use image as a base since target() can be overriden by child class
+        mask = torch.ones((1, image.shape[-2], image.shape[-1]), dtype=torch.float32)
 
         image = tv_tensors.Image(crop_image_to_mask(image, mask))
         target = tv_tensors.Mask(crop_image_to_mask(target, mask))