terg

managers/BaseManager.py

@@ -2,6 +2,7 @@
 import cv2
 import json
 import torch
+import pickle
 import datetime
 import numpy as np
 import pandas as pd
@@ -191,7 +192,9 @@ def load_data(self):
         # Create default dataloaders
         if self.config['mode'] == 'inference':
             train_df, valid_df = self.get_seg_dataframes()
+            print("train_df: ", train_df, "valid_df: ", valid_df)
             _, valid_loader = self.get_dataloaders(train_df, valid_df, 'default')
+            print("valid_loader: ", valid_loader)
             self.data_loaders = {'valid_loader': valid_loader}
             return
 
@@ -661,11 +664,18 @@ def infer(self):
 
         confusion_matrix = None
         with torch.no_grad():
+            variables_list = []
             for rec_num, (img, lbl, metadata) in enumerate(self.data_loaders['valid_loader']):
                 print("\r Inference on {}".format(rec_num), end='', flush=True)
                 img, lbl = img.to(self.device), lbl.to(self.device)
                 # noinspection PyUnboundLocalVariable
                 output = self.model(img.float()) if not self.config['tta'] else tta_model(img.float())
+                print("type of output :", type(output))
+                print("type of output :", output.shape)
+                variables_list.append(output.cpu())
+                with open('output.pkl', 'wb') as h:
+                    pickle.dump(variables_list, h)
+
                 confusion_matrix = t_get_confusion_matrix(output, lbl, confusion_matrix)
                 if rec_num in np.round(np.linspace(0, len(self.data_loaders['valid_loader']) - 1, self.max_valid_imgs)):
                     if not isinstance(self.model, Ensemble):

utils/defaults.py

@@ -1,7 +1,8 @@
 DATA_SPLITS = [  # Pre-defined splits of the videos, to be used generally
     [[1], [5]],  # Split 0: debugging
     [[1, 3, 4, 6, 8, 9, 10, 11, 13, 14, 15, 17, 18, 19, 20, 21, 23, 24, 25], [5, 7, 16, 2, 12, 22]],  # train-[val,test]
-    [[1, 3, 4, 6, 8, 9, 10, 11, 13, 14, 15, 17, 18, 19, 20, 21, 23, 24, 25], [5, 7, 16], [2, 12, 22]],  # train-val-test
+    # [[1, 3, 4, 6, 8, 9, 10, 11, 13, 14, 15, 17, 18, 19, 20, 21, 23, 24, 25], [5, 7, 16], [2, 12, 22]],  # train-val-test
+    [[3],[2],[1]],
     [list(range(1, 26)), [5, 7, 16, 2, 12, 22]],  # Split 2 (all data)
     [[1, 8, 9, 10, 14, 15, 21, 23, 24], [5, 7, 16, 2, 12, 22]],     # Split 3: "50% of data" (1729 frames, 49.3%)
     [[10, 14, 21, 24], [5, 7, 16, 2, 12, 22]],                      # Split 4: "25% of data" (834 frames, 23.8%)

utils/torch_utils.py

@@ -1,6 +1,9 @@
 import torch
 from torch.nn import Conv2d, ConvTranspose2d, ReLU, Sequential, BatchNorm2d
 from torch.nn.functional import one_hot, unfold, pad
+import pickle
+import numpy as np
+import matplotlib.pyplot as plt
 from utils import softmax, get_inverse_affine_matrix, rotate, InterpolationMode, CLASS_INFO, DEFAULT_VALUES
 
 
@@ -234,8 +237,30 @@ def t_get_confusion_matrix(prediction: torch.Tensor, target: torch.Tensor, exist
             one_hot_target = one_hot_target[:, :-1]
         else:
             one_hot_target = one_hot(t, num_classes)
+        print("one_hot_pred: ",one_hot_pred.shape)
+        print("one_hot_target: ", one_hot_target.shape)
+
+        print(type(one_hot_pred))
+
+
         confusion_matrix = torch.matmul(one_hot_pred.to(torch.float), one_hot_target.to(torch.float)).to(torch.int)
         # [C, N*H*W] x [N*H*W, C] = [C, C]
+
+        with open('one_hot_pred.pkl', 'wb') as f:
+            pickle.dump(one_hot_pred.cpu(), f)
+
+        # one_hot_encoded = one_hot_pred
+        # one_hot_encoded_np = one_hot_encoded.cpu().numpy()
+        # original_shape = (960, 540, 8)
+        # one_hot_encoded_np = one_hot_encoded_np.reshape(original_shape)
+        # segmentation_mask = np.argmax(one_hot_encoded_np, axis=-1)
+        # # Display the segmentation mask using matplotlib
+        # plt.imshow(segmentation_mask, cmap='tab20')  # 'tab20' provides a colormap with 20 distinct colors
+        # plt.colorbar()
+        # plt.title("Segmentation Mask")
+        # plt.show()
+
+
         if existing_matrix is not None:
             confusion_matrix += existing_matrix
         return confusion_matrix