Segmentation refactored, now for the point tracking

Author: Henningson

source/GUI/ImageViewerWidget.py

@@ -52,10 +52,10 @@ def convertImage(self, image):
 
         h, w, ch = image.shape
         bytesPerLine = ch * w
-        return QImage(image.data, w, h, bytesPerLine, QImage.Format_BGR888)
+        return QImage(image.copy().data, w, h, bytesPerLine, QImage.Format_BGR888)
 
     def updateImage(self, image, widget):
-        widget.setPixmap(QPixmap.fromImage(self.convertImage(image).scaledToWidth(128)))
+        widget.setPixmap(QPixmap.fromImage(self.convertImage(image)))
 
     def getWidget(self, key):
         return self.imageDICT[key]

source/GUI/MainMenuWidget.py

@@ -28,8 +28,8 @@ def __init__(self, viewer_palette, parent=None):
             "Segmentation",
             [
                 ("Koc et al", "checkbox", False),
-                ("Neural Segmentation", "checkbox", False),
-                ("Silicone Segmentation", "checkbox", True),
+                ("Neural Segmentation", "checkbox", True),
+                ("Silicone Segmentation", "checkbox", False),
             ],
         )
         self.addSubMenu(

source/KocSegmentation.py

@@ -198,9 +198,9 @@ def generateSegmentationData(self):
         self.generate()
 
         for i, image in enumerate(self.images):
-            self._segmentations.append(self.segmentImage(image))
-            self._glottal_outlines.append(self.computeGlottalOutline(i))
-            self._glottal_midlines.append(self.computeGlottalMidline(i))
+            self.segmentations.append(self.segmentImage(image))
+            self.glottal_outlines.append(self.computeGlottalOutline(i))
+            self.glottal_midlines.append(self.computeGlottalMidline(i))
 
         self.closedGlottisIndex = self.estimateClosedGlottis()
         self.openGlottisIndex = self.estimateOpenGlottis()

source/NeuralSegmentation.py

@@ -70,6 +70,125 @@ def forward(self, x):
 
 
 
+class DoubleConvB(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DoubleConvB, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True))
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Decoder(nn.Module):
+    def __init__(self, encoder, out_channels, features):
+        super(Decoder, self).__init__()
+        self.ups = nn.ModuleList()
+        self.encoder = encoder
+        self.out_channels=out_channels
+
+        for feature in reversed(features):
+            self.ups.append(nn.ConvTranspose2d(feature*2, feature, kernel_size=2, stride=2))
+            self.ups.append(DoubleConvB(feature*2, feature))
+
+
+    def forward(self, x):
+        for idx in range(0, len(self.ups), 2):
+            x = self.ups[idx](x)
+            skip_connection = self.encoder.skip_connections[idx//2]
+
+            if x.shape != skip_connection.shape:
+                x = TF.resize(x, size=skip_connection.shape[2:])
+
+            concat_skip = torch.cat((skip_connection, x), dim=1)
+            x = self.ups[idx+1](concat_skip)
+
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, in_channels, features):
+        super(Encoder, self).__init__()
+        self.downs = nn.ModuleList()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.in_channels = in_channels
+        
+        #Downsampling
+        for feature in features:
+            self.downs.append(DoubleConvB(in_channels, feature))
+            in_channels = feature
+
+    def forward(self, x):
+        self.skip_connections = []
+        for down in self.downs:
+            x = down(x)
+            self.skip_connections.append(x)
+            x = self.pool(x)
+
+        self.skip_connections = self.skip_connections[::-1]
+
+        return x
+
+
+
+class UNETNew(nn.Module):
+    def __init__(self, config={'in_channels': 3, 'out_channels': 4, 'features': [32, 64, 128, 256, 512]}, state_dict=None, pretrain=False, device="cuda"):
+        super(UNETNew, self).__init__()
+        try:
+            in_channels = config['in_channels']
+        except:
+            in_channels = 3
+
+        try: 
+            out_channels = config['out_channels']
+        except:
+            out_channels = 4
+            
+        features = config['features']
+
+        self.bottleneck_size = features[-1]*2
+
+        self.encoder = Encoder(in_channels, features)
+        self.decoder = Decoder(self.encoder, out_channels, features)
+        self.bottleneck = DoubleConv(features[-1], self.bottleneck_size)
+        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
+
+        if state_dict:
+            self.load_from_dict(state_dict)
+
+        if pretrain:
+            self.encoder.requires_grad_ = False
+
+    def get_statedict(self):
+        return {"Encoder": self.encoder.state_dict(),
+                "Bottleneck": self.bottleneck.state_dict(),
+                "Decoder": self.decoder.state_dict(),
+                "LastConv": self.final_conv.state_dict()}
+
+    def load_from_dict(self, dict):
+        self.encoder.load_state_dict(dict["Encoder"])
+        self.bottleneck.load_state_dict(dict["Bottleneck"])
+        self.decoder.load_state_dict(dict["Decoder"])
+        
+        try:
+            self.final_conv.load_state_dict(dict["LastConv"])
+        except:
+            print("Final conv not initialized.")
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.bottleneck(x)
+        x = self.decoder(x)
+
+        return self.final_conv(x)
+
+
+
 class NeuralSegmentator(BaseSegmentator):
     def __init__(self, images, path="assets/model.pth.tar"):
         super().__init__(images)

source/Viewer.py

@@ -13,6 +13,7 @@
 import helper
 import igl
 import KocSegmentation
+import kornia
 import Laser
 import Mesh
 import NeuralSegmentation
@@ -24,6 +25,7 @@
 import SiliconeSegmentation
 import SiliconeSurfaceReconstruction
 import surface_reconstruction
+import torch
 import Triangulation
 import VoronoiRHC
 from GraphWidget import GraphWidget
@@ -158,11 +160,18 @@ def __init__(self):
         self.timer_thread.start()
         self.image_timer_thread.start()
 
+        path = "/media/nu94waro/Windows_C/save/datasets/HLEDataset/dataset"
         self.loadData(
             "assets/camera_calibration.json",
             "assets/laser_calibration.json",
             "assets/example_vid.avi",
         )
+        '''
+        self.loadData(
+            os.path.join(path, "camera_calibration.json"),
+            os.path.join(path, "laser_calibration.json"),
+            os.path.join(path, "MK/MK.avi"),
+        )'''
 
         self._reconstruction_pipeline = reconstruction_pipeline.ReconstructionPipeline(
             self.camera, 
@@ -442,70 +451,36 @@ def loadData(self, camera_path, laser_path, video_path):
 
         self.images_set = True
 
+    
+
+
+
     def segmentImages(self):
-        if self.menu_widget.widget().getSubmenuValue("Segmentation", "Koc et al"):
-            self.segmentator = KocSegmentation.KocSegmentator(self.images)
-        elif self.menu_widget.widget().getSubmenuValue("Segmentation", "Neural Segmentation"):
-            self.segmentator = NeuralSegmentation.NeuralSegmentator(self.images)
+        segmentator: feature_estimation.FeatureEstimator = None
+        if self.menu_widget.widget().getSubmenuValue("Segmentation", "Neural Segmentation"):
+            segmentator = feature_estimation.NeuralFeatureEstimator("bla")
         elif self.menu_widget.widget().getSubmenuValue("Segmentation", "Silicone Segmentation"):
-            self.segmentator = SiliconeSegmentation.SiliconeSegmentator(self.images)
+            segmentator = feature_estimation.SiliconeFeatureEstimator()
         else:
             print("Please choose a Segmentation Algorithm")
 
-        x, w, y, h = self.segmentator.getROI()
-        self.roi = self.segmentator.getROIImage()
+        self._reconstruction_pipeline.set_feature_estimator(segmentator)
+        
+        images = torch.from_numpy(np.stack(self.images)).to("cuda")
+        segmentator.compute_features(images)
 
         segmentations = list()
-        laserdots = list()
-
-        for index in range(len(self.segmentator)):
-            base_image = self.segmentator.getImage(index).copy()
+        feature_images = segmentator.create_feature_images()
 
-            segmentation_image = self.segmentator.getSegmentation(index).copy()
-            gml_a, gml_b = self.segmentator.getGlottalMidline(index)
+        for feature_image in feature_images:
+            segmentations.append(feature_image.permute(1, 2, 0).detach().cpu().numpy().astype(np.uint8))
 
-            segmentation_image = cv2.cvtColor(segmentation_image, cv2.COLOR_GRAY2BGR)
-
-            cv2.rectangle(
-                segmentation_image,
-                (x, y),
-                (x + w, y + h),
-                color=(255, 0, 0),
-                thickness=2,
-            )
-            try:
-                cv2.line(
-                    segmentation_image,
-                    gml_a.astype(np.int32),
-                    gml_b.astype(np.int32),
-                    color=(125, 125, 0),
-                    thickness=2,
-                )
-            except:
-                pass
-            segmentations.append(
-                cv2.cvtColor(base_image, cv2.COLOR_GRAY2BGR) | segmentation_image
-            )
-
-            laserdot_image = self.segmentator.getLocalMaxima(index).copy()
-            laserdot_image = cv2.dilate(laserdot_image, np.ones((3, 3)))
-            laserdot_image = np.where(laserdot_image > 0, 255, 0).astype(np.uint8)
-            laserdot_image = cv2.cvtColor(laserdot_image, cv2.COLOR_GRAY2BGR)
-            laserdot_image[:, :, [0, 2]] = 0
-            laserdots.append(
-                cv2.cvtColor(base_image, cv2.COLOR_GRAY2BGR) | laserdot_image
-            )
-
-        glottal_area_waveform = [
-            len(self.segmentator.getSegmentation(index).nonzero()[0])
-            for index in range(len(self.segmentator))
-        ]
         self.graph_widget.updateGraph(
-            glottal_area_waveform, self.graph_widget.glottal_seg_graph
+            segmentator.glottalAreaWaveform().tolist(), self.graph_widget.glottal_seg_graph
         )
 
         self.segmentations = segmentations
-        self.laserdots = laserdots
+        self.laserdots = segmentations
 
     def buildCorrespondences(self):
         min_search_space = float(

source/cv.py

@@ -105,10 +105,10 @@ def compute_segmentation_outline(segmentation: torch.tensor, kernel_size=3, bord
     Returns:
         border: (B, 1, H, W) tensor of borders
     """
-    kernel = torch.ones((1, 1, kernel_size, kernel_size), device=segmentation.device)
+    kernel = torch.ones((kernel_size, kernel_size), device=segmentation.device)
 
-    dilated = kornia.morphology.dilation(segmentation, kernel)
-    eroded = kornia.morphology.erosion(segmentation, kernel)
+    dilated = kornia.morphology.dilation(segmentation.unsqueeze(0).unsqueeze(0).float(), kernel).squeeze()
+    eroded = kornia.morphology.erosion(segmentation.unsqueeze(0).unsqueeze(0).float(), kernel).squeeze()
 
     if border_type == "both":
         border = dilated - eroded
@@ -145,9 +145,9 @@ def windows_out_of_bounds(indices, image_size, pad):
 def extractWindow(batch, indices, window_size=7, device="cuda"):
     # Clean Windows, such that no image boundaries are hit
 
-    batch_index = indices[:, 0].int()
-    y = indices[:, 2].floor().int()
-    x = indices[:, 1].floor().int()
+    batch_index = indices[:, 0].long()
+    y = indices[:, 2].long()
+    x = indices[:, 1].long()
 
     y = windows_out_of_bounds(y, batch.shape[1], window_size // 2)
     x = windows_out_of_bounds(x, batch.shape[2], window_size // 2)