Rhapso 0.1.92__py3-none-any.whl

Rhapso/solver/concatenate_models.py

@@ -0,0 +1,116 @@
+from copy import deepcopy
+
+"""
+Concatenate Models stitches together the results of two alignment rounds (split-affine)
+"""
+
+class ConcatenateModels:
+    def __init__(self, tiles, tiles_round_2, groups, validation_stats, validation_stats_round_2, view_map):
+        self.tiles = tiles
+        self.tiles_round_2 = tiles_round_2
+        self.groups = groups
+        self.validation_stats = validation_stats
+        self.validation_stats_round_2 = validation_stats_round_2
+        self.view_map = view_map
+
+    def map_models_back_from_groups(self):
+        """
+        Copy solved model and matches back onto every tile in the same group
+        """
+        view_to_group = {v: gi for gi, g in enumerate(self.groups) for v in g.get('views', [])}
+
+        for rep_tile in self.tiles_round_2:
+            rep_view = rep_tile['view']
+            gi = view_to_group.get(rep_view)
+            if gi is None:
+                continue  
+
+            group_views = self.groups[gi]['views']
+            rep_model  = deepcopy(rep_tile.get('model', {}))
+            rep_matches = list(rep_tile.get('matches', []))  
+
+            # Normalize rep connections into dict
+            rep_conns = []
+            for c in rep_tile.get('connected_tiles', []):
+                v = c['view'] if isinstance(c, dict) else c
+                rep_conns.append({'view': v, 'tile': self.view_map.get(v)})
+
+            # Propagate to every member of the group
+            for v in group_views:
+                t = self.view_map.get(v)
+                if not t:
+                    continue
+
+                t['model']   = deepcopy(rep_model)     
+                t['matches'] = list(rep_matches)
+
+        return list(self.view_map.values())
+    
+    def preconcatenate_affine(self, m1, m2):
+        """
+        Compose two 3x4 affine transforms stored as dicts
+        """
+        m00,m01,m02,m03 = m1['m00'],m1['m01'],m1['m02'],m1['m03']
+        m10,m11,m12,m13 = m1['m10'],m1['m11'],m1['m12'],m1['m13']
+        m20,m21,m22,m23 = m1['m20'],m1['m21'],m1['m22'],m1['m23']
+
+        n00,n01,n02,n03 = m2['m00'],m2['m01'],m2['m02'],m2['m03']
+        n10,n11,n12,n13 = m2['m10'],m2['m11'],m2['m12'],m2['m13']
+        n20,n21,n22,n23 = m2['m20'],m2['m21'],m2['m22'],m2['m23']
+
+        return {
+            'm00': n00*m00 + n01*m10 + n02*m20,
+            'm01': n00*m01 + n01*m11 + n02*m21,
+            'm02': n00*m02 + n01*m12 + n02*m22,
+            'm03': n00*m03 + n01*m13 + n02*m23 + n03,
+
+            'm10': n10*m00 + n11*m10 + n12*m20,
+            'm11': n10*m01 + n11*m11 + n12*m21,
+            'm12': n10*m02 + n11*m12 + n12*m22,
+            'm13': n10*m03 + n11*m13 + n12*m23 + n13,
+
+            'm20': n20*m00 + n21*m10 + n22*m20,
+            'm21': n20*m01 + n21*m11 + n22*m21,
+            'm22': n20*m02 + n21*m12 + n22*m22,
+            'm23': n20*m03 + n21*m13 + n22*m23 + n23,
+        }
+    
+    def merge_validation_stats(self, v1, v2):
+        """
+        Merge two validation-metrics dicts by concatenating per-tile from v2 after v1, offsetting v2's iteration numbers 
+        and tagging entries
+        """
+        out = deepcopy(v1) if v1 else {}
+        s1 = out.setdefault('solver_metrics_per_tile', {}).setdefault('stats', [])
+        s2 = (v2 or {}).get('solver_metrics_per_tile', {}).get('stats', []) or []
+        
+        # tag round and offset iterations
+        if s1 and 'round' not in s1[0]:
+            for x in s1: x['round'] = 1
+        offset = (s1[-1]['iteration'] + 1) if s1 else 0
+        
+        for x in s2:
+            y = dict(x)
+            y['iteration'] = x.get('iteration', 0) + offset
+            y['round'] = 2
+            s1.append(y)
+            
+        return out
+
+    def run(self):
+        """
+        Executes the entry point of the script.
+        """
+        view_map = self.map_models_back_from_groups()
+        combined_validation_stats = self.merge_validation_stats(self.validation_stats, self.validation_stats_round_2)
+
+        tiles_round_1 = {t['view']: t for t in self.tiles}
+        tiles_round_2  = {t['view']: t for t in view_map}
+
+        for vid in tiles_round_1.keys() & tiles_round_2.keys():
+            m1 = tiles_round_1[vid]['model']['regularized']
+            m2 = tiles_round_2[vid]['model']['regularized']
+            tiles_round_2[vid]['model']['regularized'] = self.preconcatenate_affine(m1, m2) 
+        
+        tiles_round_2 = [tiles_round_2[t['view']] for t in view_map]
+        return tiles_round_2, combined_validation_stats

Rhapso/solver/connected_graphs.py

@@ -0,0 +1,111 @@
+
+"""
+Connected Graphs splits a tile set into connected components using each tiles connected_tiles links (split-affine).
+"""
+
+class ConnectedGraphs:
+    def __init__(self, tiles, dataframes):
+        self.tiles = tiles
+        self.view_registrations = dataframes['view_registrations']
+
+    def identify_connected_graphs(self):
+        """
+        Build an undirected graph of tiles
+        """
+        view_to_tile = {t['view']: t for t in self.tiles}
+
+        # undirected adjacency by view_id
+        adj = {v: set() for v in view_to_tile}
+        for t in self.tiles:
+            v = t['view']
+            for conn in t.get('connected_tiles', []):
+                nv = conn['view'] if isinstance(conn, dict) else conn
+                if nv in view_to_tile and nv != v:
+                    adj[v].add(nv)
+                    adj[nv].add(v)
+
+        graphs, visited = [], set()
+        for v in adj:
+            if v in visited:
+                continue
+            stack = [v]
+            comp_views = set([v])
+            visited.add(v)
+            while stack:
+                u = stack.pop()
+                for w in adj[u]:
+                    if w not in visited:
+                        visited.add(w)
+                        comp_views.add(w)
+                        stack.append(w)
+
+            comp = sorted(comp_views) 
+            graphs.append([view_to_tile[x] for x in comp])
+
+        return graphs
+
+    def assemble_views(self, connected):
+        """
+        Normalize a collection of connected tiles to a set of view IDs
+        """
+        if connected and isinstance(next(iter(connected)), dict):
+            connected_view_ids = {t['view'] for t in connected}
+        else:
+            connected_view_ids = set(connected)
+
+        group_views = [t['view'] for t in self.tiles if t['view'] in connected_view_ids]
+
+        return {'views': group_views}
+
+    def create_wlpmc(self):
+        """
+        Initialize weak link point match correspondences
+        """
+        views = [t['view'] for t in self.tiles]
+        return {
+            'view': views,
+            'models': self.tiles,
+            'overlap_detection': None,
+            'view_registrations': self.view_registrations
+        }
+    
+    def label_subsets(self, group):
+        """
+        Sort a group's view IDs by setup, then collapse consecutive setup numbers into ranges
+        """
+        vs = sorted(group['views'], key=lambda v: int(v.partition('setup:')[2].split()[0]))
+        tp = int(vs[0].partition('timepoint:')[2].split(',')[0]) if vs else 0
+
+        nums = sorted({int(v.partition('setup:')[2].split()[0]) for v in vs})
+        labels, starts = [], []
+        if nums:
+            s = e = nums[0]
+            for x in nums[1:]:
+                if x == e + 1:
+                    e = x
+                else:
+                    labels.append(f"{tp}-{s} >-> {tp}-{e}"); starts.append(s)
+                    s = e = x
+            labels.append(f"{tp}-{s} >-> {tp}-{e}"); starts.append(s)
+
+        group['views'] = vs
+        group['subset_labels'] = labels
+        return group
+
+    def run(self):
+        """
+        Executes the entry point of the script.
+        """
+        graph_sets = self.identify_connected_graphs()
+        groups_new = []
+
+        if len(graph_sets) == 1:
+            return self.tiles, groups_new
+        else:
+            for connected in graph_sets:
+                group = self.assemble_views(connected)
+                group = self.label_subsets(group)
+                groups_new.append(group)
+        
+        wlpmc = self.create_wlpmc()
+        return wlpmc, groups_new

Rhapso/solver/data_prep.py

@@ -0,0 +1,181 @@
+import zarr
+import json
+import os
+import s3fs
+
+"""
+Data Prep fetches and preps n5 interest points data
+"""
+
+class DataPrep():
+    def __init__(self, interest_points_df, view_transform_matrices, xml_file_path, n5_input_path):
+        self.interest_points_df = interest_points_df
+        self.view_transform_matrices = view_transform_matrices
+        self.xml_file_path = xml_file_path
+        self.n5_input_path = n5_input_path
+
+        self.connected_views = {} 
+        self.corresponding_interest_points = {}
+        self.interest_points = {}
+        self.label_map_global = {}
+    
+    def get_connected_views_from_n5(self):
+        """
+        Loads connected view mappings from N5 metadata, supporting both S3 and local sources.
+        """
+        for _, row in self.interest_points_df.iterrows():
+            view_id = f"timepoint: {row['timepoint']}, setup: {row['setup']}"
+
+            if self.n5_input_path.startswith("s3://"):
+                s3 = s3fs.S3FileSystem(anon=False)
+                root = self.n5_input_path.replace("s3://", "", 1)
+                path = root + "interestpoints.n5"
+                store = s3fs.S3Map(root=path, s3=s3)
+                root = zarr.open(store, mode='r')
+                correspondences_key = f"{row['path']}/correspondences"
+                try:
+                    self.connected_views[view_id] = root[correspondences_key].attrs["idMap"]
+                except:
+                    print(f"No connected views for tile {view_id}")
+
+            else:
+                n5_root = os.path.join(self.n5_input_path, "interestpoints.n5")
+                store = zarr.N5Store(n5_root)
+                root = zarr.open(store, mode="r")
+                correspondences_key = f"{row['path']}/correspondences"
+                try:
+                    self.connected_views[view_id] = self.load_json_data(correspondences_key)
+                except:
+                    print(f"No connected views for tile {view_id}")            
+    
+    def load_json_data(self, json_path):
+        try:
+            path = self.n5_input_path + json_path
+            if not os.path.exists(path):
+                return {}
+            with open(path, 'r') as f:
+                obj = json.load(f)
+            id_map = obj.get('idMap', {})
+            return id_map if isinstance(id_map, dict) else {}
+        except Exception:
+            return {}
+        
+    def get_corresponding_data_from_n5(self):
+        """
+        Parses and transforms corresponding interest point data from N5 format into world space coordinates.
+        """
+        if self.n5_input_path.startswith("s3://"):
+            s3 = s3fs.S3FileSystem(anon=False)
+            root = self.n5_input_path.replace("s3://", "", 1)
+            path = root + "interestpoints.n5"
+            store = s3fs.S3Map(root=path, s3=s3)
+        else:
+            path = self.n5_input_path + "interestpoints.n5"
+            store = zarr.N5Store(path)
+
+        root = zarr.open(store, mode='r')
+
+        for _, row in self.interest_points_df.iterrows():
+            view_id = f"timepoint: {row['timepoint']}, setup: {row['setup']}"  
+            correspondences_prefix = f"{row['path']}/correspondences"
+            attributes_path = 'interestpoints.n5' + f"/{row['path']}/correspondences/attributes.json"
+
+            # Load JSON data for idMap
+            if self.n5_input_path.startswith("s3://"):
+                try:
+                    id_map = root[correspondences_prefix].attrs['idMap']
+                except Exception:
+                    continue
+            else:
+                id_map = self.load_json_data(attributes_path)
+                if not id_map:
+                    continue
+            
+            try:
+                interest_points_index_map = root[correspondences_prefix + '/data'][:]
+            except (KeyError, FileNotFoundError, AttributeError, TypeError):
+                print(f"⚠️ Skipping {view_id}: missing correspondences.")
+                continue
+
+            # Load corresponding interest points data
+            for ip_index, corr_index, corr_group_id in interest_points_index_map:
+                if corr_group_id == view_id:
+                    continue
+
+                corresponding_view_id = next((k for k, v in id_map.items() if v == int(corr_group_id)), None)
+                parts = corresponding_view_id.split(',')
+                timepoint, setup, label = parts[0], parts[1], parts[2]
+                corresponding_view_id = f"timepoint: {timepoint}, setup: {setup}"
+
+                ip = self.interest_points[view_id][label][int(ip_index)]
+                corr_ip = self.interest_points[corresponding_view_id][label][int(corr_index)]
+
+                if view_id not in self.corresponding_interest_points:
+                    self.corresponding_interest_points[view_id] = [] 
+                
+                self.corresponding_interest_points[view_id].append({
+                    "detection_id": ip_index,
+                    "detection_p1": ip,
+                    "corresponding_detection_id":  corr_index,
+                    "corresponding_detection_p2": corr_ip,
+                    "corresponding_view_id": corresponding_view_id,
+                    "label": label
+                })
+    
+    def get_all_interest_points_from_n5(self):
+        """
+        Loads raw interest point coordinates from N5 storage into memory, keyed by view ID.
+        """
+        if self.n5_input_path.startswith("s3://"):
+            s3 = s3fs.S3FileSystem(anon=False)
+            root = self.n5_input_path.replace("s3://", "", 1)
+            path = root + "interestpoints.n5"
+            store = s3fs.S3Map(root=path, s3=s3)
+        else:
+            path = self.n5_input_path + "interestpoints.n5"
+            store = zarr.N5Store(path)
+
+        root = zarr.open(store, mode='r')
+
+        for _, row in self.interest_points_df.iterrows():
+            view_id = f"timepoint: {row['timepoint']}, setup: {row['setup']}"  
+            interestpoints_prefix = f"{row['path']}/interestpoints/loc/"
+            interest_points = root[interestpoints_prefix][:]
+            # interest_points = root[interestpoints_prefix][:] if interestpoints_prefix in root else np.empty((0, 3), dtype=np.float32)
+            label = str(row['path']).replace('\\','/').lstrip('/').split('/', 2)[1]
+            self.interest_points.setdefault(view_id, {})[label] = interest_points
+                            
+    def build_label_map(self):
+        """
+        Constructs a mapping of labels for each view ID from the interest points dataframe.
+        """
+        for _, row in self.interest_points_df.iterrows():
+            view_id_key = f"timepoint: {row['timepoint']}, setup: {row['setup']}"
+            
+            if view_id_key not in self.label_map_global:
+                self.label_map_global[view_id_key] = {}
+
+            self.label_map_global[view_id_key][row['label']] = 1.0  
+    
+    def run(self):
+        """
+        Executes the entry point of the script.
+        """
+        self.build_label_map()
+        self.get_all_interest_points_from_n5()
+        self.get_corresponding_data_from_n5()
+        self.get_connected_views_from_n5()
+
+        view_id_set = set()
+        for k in self.corresponding_interest_points.keys():
+            try:
+                parts = [p.strip() for p in k.split(',')]
+                tp = parts[0].split(':')[-1].strip()
+                su = parts[1].split(':')[-1].strip()
+                view_id_set.add((str(tp), str(su)))
+            except Exception:
+                continue
+
+        self.view_id_set = sorted(view_id_set, key=lambda x: (int(x[0]), int(x[1])))
+
+        return self.connected_views, self.corresponding_interest_points, self.interest_points, self.label_map_global, self.view_id_set