fix output number

Author: taigw

pymic/net_run/infer_func.py

@@ -21,6 +21,7 @@ def volume_infer(image, net, device, class_num,
         outputs = net(image)
         if(isinstance(outputs, tuple) or isinstance(outputs, list)):
             outputs = [item.cpu().numpy() for item in outputs]
+            outputs = outputs[:output_num]
         else:
             outputs = outputs.cpu().numpy()
     else:
@@ -101,6 +102,8 @@ def volume_infer_by_patch(image, net, device, class_num,
         if(not(isinstance(out_mini_batch, tuple) or isinstance(out_mini_batch, list))):
             out_mini_batch = [out_mini_batch]
         out_mini_batch  = [item.cpu().numpy() for item in out_mini_batch]
+
+        # use a mask to store overlapping regions
         mask_mini_batch = np.ones_like(out_mini_batch[0])
         for batch_idx in range(batch_start_idx, batch_end_idx):
             crop_start = sub_image_starts[batch_idx]

pymic/net_run/net_run_agent.py

@@ -370,6 +370,11 @@ def test_time_dropout(m):
 
                 prob_list = [scipy.special.softmax(predict[0], axis = 0) for predict in predict_list]
                 if(multi_pred_avg):
+                    if(output_num == 1):
+                        raise ValueError("multiple predictions expected, but output_num was set to 1")
+                    if(output_num != len(prob_list)):
+                        raise ValueError("expected output_num was set to {0:}, but {1:} outputs obtained".format(
+                              output_dir, len(prob_list)))
                     prob_stack   = np.asarray(prob_list, np.float32)
                     prob   = np.mean(prob_stack, axis = 0)
                     var    = np.var(prob_stack, axis = 0)