Merge branch 'master' into python3.8

3 files changed, 35 insertions, 8 deletions

diff --git a/config.py b/config.py
index 8abeba3..c928067 100644
--- a/config.py
+++ b/config.py
@@ -94,9 +94,9 @@ config: Configuration = {
             'final_gamma': 0.01,
 
             # Local parameters (override global ones)
-            'hpm': {
-                'final_gamma': 0.001
-            }
+            # 'hpm': {
+            #     'final_gamma': 0.001
+            # }
         }
     },
     # Model metadata
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index 6e72e38..fba3485 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -59,7 +59,7 @@ class RGBPartNet(nn.Module):
         if self.training:
             return x_c, x_p, ae_losses, images
         else:
-            return torch.cat((x_c, x_p)).unsqueeze(1).view(-1)
+            return torch.cat((x_c.view(-1), x_p.view(-1)))
 
     def _disentangle(self, x_c1_t2, x_c2_t2=None):
         n, t, c, h, w = x_c1_t2.size()
diff --git a/utils/sampler.py b/utils/sampler.py
index 0977f94..581d7a2 100644
--- a/utils/sampler.py
+++ b/utils/sampler.py
@@ -15,7 +15,18 @@ class TripletSampler(data.Sampler):
     ):
         super().__init__(data_source)
         self.metadata_labels = data_source.metadata['labels']
+        metadata_conditions = data_source.metadata['conditions']
+        self.subsets = {}
+        for condition in metadata_conditions:
+            pre, _ = condition.split('-')
+            if self.subsets.get(pre, None) is None:
+                self.subsets[pre] = []
+            self.subsets[pre].append(condition)
+        self.num_subsets = len(self.subsets)
+        self.num_seq = {pre: len(seq) for (pre, seq) in self.subsets.items()}
+        self.min_num_seq = min(self.num_seq.values())
         self.labels = data_source.labels
+        self.conditions = data_source.conditions
         self.length = len(self.labels)
         self.indexes = np.arange(0, self.length)
         (self.pr, self.k) = batch_size
@@ -26,15 +37,31 @@ class TripletSampler(data.Sampler):
             # Sample pr subjects by sampling labels appeared in dataset
             sampled_subjects = random.sample(self.metadata_labels, k=self.pr)
             for label in sampled_subjects:
-                clips_from_subject = self.indexes[self.labels == label].tolist()
+                mask = self.labels == label
+                # Fix unbalanced datasets
+                if self.num_subsets > 1:
+                    condition_mask = np.zeros(self.conditions.shape, dtype=bool)
+                    for num, conditions_ in zip(
+                            self.num_seq.values(), self.subsets.values()
+                    ):
+                        if num > self.min_num_seq:
+                            conditions = random.sample(
+                                conditions_, self.min_num_seq
+                            )
+                        else:
+                            conditions = conditions_
+                        for condition in conditions:
+                            condition_mask |= self.conditions == condition
+                    mask &= condition_mask
+                clips = self.indexes[mask].tolist()
                 # Sample k clips from the subject without replacement if
                 # have enough clips, k more clips will sampled for
                 # disentanglement
                 k = self.k * 2
-                if len(clips_from_subject) >= k:
-                    _sampled_indexes = random.sample(clips_from_subject, k=k)
+                if len(clips) >= k:
+                    _sampled_indexes = random.sample(clips, k=k)
                 else:
-                    _sampled_indexes = random.choices(clips_from_subject, k=k)
+                    _sampled_indexes = random.choices(clips, k=k)
                 sampled_indexes += _sampled_indexes
 
             yield sampled_indexes