Fix multi-class heatmap plotting overview

Avic3nna · web-flow · commit 37f021ff1be5 · 2024-03-26T11:12:08.000+01:00
diff --git a/stamp/heatmaps/__main__.py b/stamp/heatmaps/__main__.py
@@ -168,7 +168,7 @@ def main(
             learn.model(feats.unsqueeze(-2), torch.ones((len(feats)))), dim=1
         )
         scores_2d = vals_to_im(scores, coords // stride).detach()
-        fig, axs = plt.subplots(nrows=2, ncols=min(2, len(categories)), figsize=(12, 8))
+        fig, axs = plt.subplots(nrows=2, ncols=max(2, len(categories)), figsize=(12, 8))
 
         show_class_map(
             class_ax=axs[0, 1],

Original file line number	Diff line number	Diff line change
`@@ -168,7 +168,7 @@ def main(`
`168`	`168`	`learn.model(feats.unsqueeze(-2), torch.ones((len(feats)))), dim=1`
`169`	`169`	`)`
`170`	`170`	`scores_2d = vals_to_im(scores, coords // stride).detach()`
`171`		`- fig, axs = plt.subplots(nrows=2, ncols=min(2, len(categories)), figsize=(12, 8))`
	`171`	`+ fig, axs = plt.subplots(nrows=2, ncols=max(2, len(categories)), figsize=(12, 8))`
`172`	`172`
`173`	`173`	`show_class_map(`
`174`	`174`	`class_ax=axs[0, 1],`