updates
This commit is contained in:
		
							parent
							
								
									eb6a4b5b84
								
							
						
					
					
						commit
						9a27339e04
					
				|  | @ -349,7 +349,8 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4): | |||
| 
 | ||||
|         class_prob, class_pred = torch.max(F.softmax(pred[:, 5:], 1), 1) | ||||
| 
 | ||||
|         v = ((pred[:, 4] > conf_thres) & (class_prob > .4))  # TODO examine arbitrary 0.4 thres here | ||||
|         # v = ((pred[:, 4] > conf_thres) & (class_prob > .4))  # TODO examine arbitrary 0.4 thres here | ||||
|         v = pred[:, 4] > conf_thres | ||||
|         v = v.nonzero().squeeze() | ||||
|         if len(v.shape) == 0: | ||||
|             v = v.unsqueeze(0) | ||||
|  |  | |||
		Loading…
	
		Reference in New Issue