diff --git a/topaz/algorithms.py b/topaz/algorithms.py index afbfb60..64da836 100644 --- a/topaz/algorithms.py +++ b/topaz/algorithms.py @@ -2,9 +2,9 @@ import numpy as np from scipy.optimize import linear_sum_assignment +from typing import Tuple - -def match_coordinates(targets:np.ndarray, preds:np.ndarray, radius:float) -> tuple[np.ndarray, np.ndarray]: +def match_coordinates(targets:np.ndarray, preds:np.ndarray, radius:float) -> Tuple[np.ndarray, np.ndarray]: d2 = np.sum((preds[:,np.newaxis] - targets[np.newaxis])**2, 2) cost = d2 - radius*radius cost[cost > 0] = 0 @@ -22,7 +22,7 @@ def match_coordinates(targets:np.ndarray, preds:np.ndarray, radius:float) -> tu return assignment, dist -def non_maximum_suppression(x, r:int, threshold:float=-np.inf) -> tuple[np.ndarray, np.ndarray]: +def non_maximum_suppression(x, r:int, threshold:float=-np.inf) -> Tuple[np.ndarray, np.ndarray]: ## enumerate coordinate deltas within radius/distance r width = r ii,jj = np.meshgrid(np.arange(-width,width+1), np.arange(-width,width+1)) @@ -63,7 +63,7 @@ def non_maximum_suppression(x, r:int, threshold:float=-np.inf) -> tuple[np.ndarr return scores[:j], coords[:j] -def non_maximum_suppression_3d(x:np.ndarray, r:int, scale:float=1.0, threshold:float=-np.inf) -> tuple[np.ndarray, np.ndarray]: +def non_maximum_suppression_3d(x:np.ndarray, r:int, scale:float=1.0, threshold:float=-np.inf) -> Tuple[np.ndarray, np.ndarray]: ## enumerate coordinate deltas within (possibly scaled) radius/distance r r = scale*r width = int(np.ceil(r)) diff --git a/topaz/predict.py b/topaz/predict.py index dba9671..bdb5626 100644 --- a/topaz/predict.py +++ b/topaz/predict.py @@ -2,7 +2,7 @@ import torch import numpy as np -from typing import Iterator, Iterable +from typing import Iterator, Iterable, List def batches(X:Iterable[np.ndarray], batch_size:int=1) -> Iterator[torch.Tensor]: batch = [] @@ -28,7 +28,7 @@ def score_stream(model:torch.nn.Module, images:Iterable[np.ndarray], use_cuda:bo yield logits[i] -def score(model:torch.nn.Module, images:Iterable[np.ndarray], use_cuda:bool=False, batch_size:int=1) -> list[np.ndarray]: +def score(model:torch.nn.Module, images:Iterable[np.ndarray], use_cuda:bool=False, batch_size:int=1) -> List[np.ndarray]: scores = [] for y in score_stream(model, images, use_cuda=use_cuda, batch_size=batch_size): scores.append(y)