the issue image is now being processed

plotdigitizer/image.py

@@ -2,6 +2,7 @@
 
 
 from pathlib import Path
+from collections import defaultdict
 import typing as T
 from loguru import logger
 
@@ -13,7 +14,6 @@
 from plotdigitizer import geometry
 from plotdigitizer import grid
 from plotdigitizer import plot
-from plotdigitizer.trajectory import find_trajectory
 
 
 def click_points(event, x, y, _flags, params):
@@ -54,11 +54,45 @@ def invert_image(self):
         inv_img = cv.bitwise_not(self._last())
         self._append("inverted", inv_img)
 
-    def trajectories(self):
-        return process_image(self._last(), "_image")
-
-    def extract_trajectories(self):
-        logger.info(f"Extracting trajectories from {infile}")
+    def trajectories(self, cache_key: T.Optional[str] = None):
+        img = self._last()
+        common.params_ = compute_foregrond_background_stats(img)
+        T = transform_axis(img, self.coordinates, self.indices, erase_near_axis=3)
+        assert img.std() > 0.0, "No data in the image!"
+        logger.info(f" {img.mean()}  {img.std()}")
+        if cache_key is not None:
+            save_img_in_cache(img, f"{cache_key}.transformed_axis.png")
+
+        # extract the plot that has color which is farthest from the background.
+        trajcolor = common.params_["timeseries_colors"][0]
+        img = normalize(img)
+        traj, img = self.find_trajectory(img, trajcolor, T)
+        if cache_key is not None:
+            save_img_in_cache(img, f"{cache_key}.final.png")
+        return traj
+
+
+    def find_trajectory(self, img: np.ndarray, pixel: int, T):
+        logger.info(f"Extracting trajectory for color {pixel}")
+        assert (
+            img.min() <= pixel <= img.max()
+        ), f"{pixel} is outside the range: [{img.min()}, {img.max()}]"
+
+        # Find all pixels which belongs to a trajectory.
+        o = 6
+        _clower, _cupper = _valid_px(pixel - o // 2), _valid_px(pixel + o // 2)
+
+        Y, X = np.where((img >= _clower) & (img <= _cupper))
+        traj = defaultdict(list)
+        for x, y in zip(X, Y):
+            traj[x].append(y)
+
+        assert traj, "Empty trajectory"
+
+        # this is a simple fit using median.
+        new = np.zeros_like(img)
+        res = fit_trajectory_using_median(traj, T, new)
+        return res, np.vstack((img, new))
 
     def map_axis(self):
         logger.info("Mapping axis...")
@@ -85,25 +119,6 @@ def _append(self, operation: str, img):
         self.imgs.append((operation, img))
 
 
-def process_image(img, cache_key: T.Optional[str] = None):
-    global params_
-    common.params_ = compute_foregrond_background_stats(img)
-
-    T = transform_axis(img, erase_near_axis=3)
-    assert img.std() > 0.0, "No data in the image!"
-    logger.info(f" {img.mean()}  {img.std()}")
-    if cache_key is not None:
-        save_img_in_cache(img, f"{cache_key}.transformed_axis.png")
-
-    # extract the plot that has color which is farthest from the background.
-    trajcolor = common.params_["timeseries_colors"][0]
-    img = normalize(img)
-    traj, img = find_trajectory(img, trajcolor, T)
-    if cache_key is not None:
-        save_img_in_cache(img, f"{cache_key}.final.png")
-    return traj
-
-
 def compute_foregrond_background_stats(img) -> T.Dict[str, T.Any]:
     """Compute foreground and background color."""
     params: T.Dict[str, T.Any] = {}
@@ -170,15 +185,13 @@ def axis_transformation(p, P: T.List[geometry.Point]):
     return ((sX, sY), (offX, offY))
 
 
-def transform_axis(img, erase_near_axis: int = 0):
-    global locations_
-    global points_
+def transform_axis(img, coordinates, indices, erase_near_axis: int = 0):
     # extra: extra rows and cols to erase. Help in containing error near axis.
     # compute the transformation between old and new axis.
-    T = axis_transformation(common.points_, common.locations_)
-    p = geometry.find_origin(common.locations_)
+    T = axis_transformation(indices, coordinates)
+    p = geometry.find_origin(coordinates)
     offCols, offRows = p.x, p.y
-    logger.info(f"{common.locations_} → origin {offCols}, {offRows}")
+    logger.info(f"{coordinates} → origin {offCols}, {offRows}")
     img[:, : offCols + erase_near_axis] = common.params_["background"]
     img[-offRows - erase_near_axis :, :] = common.params_["background"]
     logger.debug(f"Tranformation params: {T}")
@@ -223,3 +236,45 @@ def ask_user_to_locate_points(points, img) -> list:
             break
     logger.info("You clicked %s" % common.locations_)
     return common.locations_
+
+def _valid_px(val: int) -> int:
+    return min(max(0, val), 255)
+
+def _find_center(vec):
+    return np.median(vec)
+
+
+def fit_trajectory_using_median(traj, T, img):
+    (sX, sY), (offX, offY) = T
+    res = []
+    r, _ = img.shape
+
+    # x, y = zip(*sorted(traj.items()))
+    # logger.info((xvec, ys))
+
+    for k in sorted(traj):
+        x = k
+
+        vals = np.array(traj[k])
+
+        # For each x, we may multiple pixels in column of the image which might
+        # be y. Usually experience is that the trajectories are close to the
+        # top rather to the bottom. So we discard call pixel which are below
+        # the center of mass (median here)
+        # These are opencv pixles. So there valus starts from the top. 0
+        # belogs to top row. Therefore > rather than <.
+        avg = np.median(vals)
+        vals = vals[np.where(vals >= avg)]
+        if len(vals) == 0:
+            continue
+
+        # Still we have multiple candidates for y for each x.
+        # We find the center of these points and call it the y for given x.
+        y = _find_center(vals)
+        cv.circle(img, (x, int(y)), 1, 255, -1)
+        x1 = (x - offX) / sX
+        y1 = (r - y - offY) / sY
+        res.append((x1, y1))
+
+    # sort by x-axis.
+    return sorted(res)