Complete implementation of sex and tissue site parameters

- Add UI dropdowns for Sex and Tissue Site selection
- Update settings dataframe to include Sex and Tissue Site columns
- Pass sex and tissue_site parameters through the analysis pipeline
- Convert sex string to index (0=Male, 1=Female, None=Unknown)
- Convert tissue site string to index using tissue site mapping
- Add get_tissue_sites() helper function in ui/utils.py
- Update all inference pipeline functions to handle these parameters
- Pass sex and tissue_site_idx to Aeon model for improved inference

src/mosaic/analysis.py

@@ -154,13 +154,15 @@ def _extract_optimus_features(filtered_coords, slide_path, attrs, num_workers):
     return features
 
 
-def _run_aeon_inference(features, site_type, num_workers):
+def _run_aeon_inference(features, site_type, num_workers, sex=None, tissue_site_idx=None):
     """Run Aeon cancer subtype inference on GPU.
     
     Args:
         features: Optimus features
         site_type: Site type ("Primary" or "Metastatic")
         num_workers: Number of worker processes
+        sex: Patient sex (0=Male, 1=Female), optional
+        tissue_site_idx: Tissue site index (0-56), optional
         
     Returns:
         Aeon results DataFrame
@@ -183,6 +185,8 @@ def _run_aeon_inference(features, site_type, num_workers):
         metastatic=(site_type == "Metastatic"),
         batch_size=8,
         num_workers=num_workers,
+        sex=sex,
+        tissue_site_idx=tissue_site_idx,
         use_cpu=False,
     )
     end_time = pd.Timestamp.now()
@@ -260,6 +264,8 @@ def _run_inference_pipeline_free(
     slide_path,
     attrs,
     site_type,
+    sex,
+    tissue_site_idx,
     cancer_subtype,
     cancer_subtype_name_map,
     num_workers,
@@ -267,8 +273,8 @@ def _run_inference_pipeline_free(
 ):
     """Run inference pipeline with 60s GPU limit (for free users)."""
     return _run_inference_pipeline_impl(
-        coords, slide_path, attrs, site_type, cancer_subtype,
-        cancer_subtype_name_map, num_workers, progress
+        coords, slide_path, attrs, site_type, sex, tissue_site_idx,
+        cancer_subtype, cancer_subtype_name_map, num_workers, progress
     )
 
 
@@ -278,6 +284,8 @@ def _run_inference_pipeline_pro(
     slide_path,
     attrs,
     site_type,
+    sex,
+    tissue_site_idx,
     cancer_subtype,
     cancer_subtype_name_map,
     num_workers,
@@ -285,8 +293,8 @@ def _run_inference_pipeline_pro(
 ):
     """Run inference pipeline with 300s GPU limit (for PRO users)."""
     return _run_inference_pipeline_impl(
-        coords, slide_path, attrs, site_type, cancer_subtype,
-        cancer_subtype_name_map, num_workers, progress
+        coords, slide_path, attrs, site_type, sex, tissue_site_idx,
+        cancer_subtype, cancer_subtype_name_map, num_workers, progress
     )
 
 
@@ -295,6 +303,8 @@ def _run_inference_pipeline_impl(
     slide_path,
     attrs,
     site_type,
+    sex,
+    tissue_site_idx,
     cancer_subtype,
     cancer_subtype_name_map,
     num_workers,
@@ -351,7 +361,7 @@ def _run_inference_pipeline_impl(
     # Step 5: Run Aeon to predict histology if not supplied
     if cancer_subtype == "Unknown":
         progress(0.9, desc="Running Aeon for cancer subtype inference")
-        aeon_results = _run_aeon_inference(features, site_type, num_workers)
+        aeon_results = _run_aeon_inference(features, site_type, num_workers, sex, tissue_site_idx)
     else:
         cancer_subtype_code = cancer_subtype_name_map.get(cancer_subtype)
         aeon_results = pd.DataFrame(
@@ -509,6 +519,19 @@ def analyze_slide(
             import traceback
             logger.warning(traceback.format_exc())
     
+    # Convert sex and tissue_site to indices for Aeon model
+    sex_idx = None
+    if sex and sex != "Unknown":
+        sex_idx = 0 if sex == "Male" else 1
+    
+    tissue_site_idx = None
+    if tissue_site and tissue_site != "Unknown":
+        from mosaic.inference.data import get_tissue_site_map
+        tissue_site_map = get_tissue_site_map()
+        tissue_site_idx = tissue_site_map.get(tissue_site)
+        if tissue_site_idx is None:
+            logger.warning(f"Unknown tissue site: {tissue_site}")
+    
     if is_logged_in:
         logger.info("Using 300s GPU allocation (logged-in user)")
         aeon_results, paladin_results = _run_inference_pipeline_pro(
@@ -516,6 +539,8 @@ def analyze_slide(
             slide_path,
             attrs,
             site_type,
+            sex_idx,
+            tissue_site_idx,
             cancer_subtype,
             cancer_subtype_name_map,
             num_workers,
@@ -528,6 +553,8 @@ def analyze_slide(
             slide_path,
             attrs,
             site_type,
+            sex_idx,
+            tissue_site_idx,
             cancer_subtype,
             cancer_subtype_name_map,
             num_workers,

src/mosaic/inference/aeon.py

@@ -53,7 +53,8 @@ NUM_WORKERS = 8
 
 
 def run(
-    features, model_path, metastatic=False, batch_size=8, num_workers=8, use_cpu=False
+    features, model_path, metastatic=False, batch_size=8, num_workers=8, use_cpu=False,
+    sex=None, tissue_site_idx=None
 ):
     """Run Aeon model inference for cancer subtype prediction.
     
@@ -64,6 +65,8 @@ def run(
         batch_size: Batch size for inference
         num_workers: Number of workers for data loading
         use_cpu: Force CPU usage instead of GPU
+        sex: Patient sex (0=Male, 1=Female), optional
+        tissue_site_idx: Tissue site index (0-56), optional
         
     Returns:
         tuple: (results_df, part_embedding)
@@ -85,6 +88,8 @@ def run(
     dataset = TileFeatureTensorDataset(
         site_type=site_type,
         tile_features=features,
+        sex=sex,
+        tissue_site_idx=tissue_site_idx,
         n_max_tiles=20000,
     )
     dataloader = DataLoader(
@@ -95,6 +100,10 @@ def run(
     batch = next(iter(dataloader))
     with torch.no_grad():
         batch["tile_tensor"] = batch["tile_tensor"].to(device)
+        if "SEX" in batch:
+            batch["SEX"] = batch["SEX"].to(device)
+        if "TISSUE_SITE" in batch:
+            batch["TISSUE_SITE"] = batch["TISSUE_SITE"].to(device)
         y = model(batch)
         y["logits"][:, col_indices_to_drop] = -1e6
 

src/mosaic/inference/data.py

@@ -287,6 +287,8 @@ class TileFeatureTensorDataset(Dataset):
         self,
         site_type: SiteType,
         tile_features: np.ndarray,
+        sex: int = None,
+        tissue_site_idx: int = None,
         n_max_tiles: int = 20000,
     ) -> None:
         """Initialize the dataset.
@@ -294,12 +296,16 @@ class TileFeatureTensorDataset(Dataset):
         Args:
             site_type: the site type as str, either "Primary" or "Metastasis"
             tile_features: the tile feature array
+            sex: patient sex (0=Male, 1=Female), optional for Aeon
+            tissue_site_idx: tissue site index (0-56), optional for Aeon
             n_max_tiles: the maximum number of tiles to use as int
 
         Returns:
             None
         """
         self.site_type = site_type
+        self.sex = sex
+        self.tissue_site_idx = tissue_site_idx
         self.n_max_tiles = n_max_tiles
         self.features = self._get_features(tile_features)
 
@@ -340,7 +346,22 @@ class TileFeatureTensorDataset(Dataset):
         Returns:
             dict: the item
         """
-        return {
+        result = {
             "site": self.site_type.value,
             "tile_tensor": self.features
         }
+        
+        # Add sex and tissue_site if provided (for Aeon)
+        if self.sex is not None:
+            result["SEX"] = torch.tensor(
+                tissue_site_to_one_hot(self.sex, num_classes=3),
+                dtype=torch.float32
+            )
+        
+        if self.tissue_site_idx is not None:
+            result["TISSUE_SITE"] = torch.tensor(
+                tissue_site_to_one_hot(self.tissue_site_idx, num_classes=57),
+                dtype=torch.float32
+            )
+        
+        return result

src/mosaic/ui/app.py

@@ -18,7 +18,9 @@ from mosaic.ui.utils import (
     create_user_directory,
     load_settings,
     validate_settings,
+    get_tissue_sites,
     IHC_SUBTYPES,
+    SEX_OPTIONS,
     SETTINGS_COLUMNS,
 )
 from mosaic.analysis import analyze_slide
@@ -80,6 +82,8 @@ def analyze_slides(
             slides[idx],
             row["Segmentation Config"],
             row["Site Type"],
+            row["Sex"],
+            row["Tissue Site"],
             row["Cancer Subtype"],
             cancer_subtype_name_map,
             row["IHC Subtype"],
@@ -177,6 +181,16 @@ def launch_gradio(server_name, server_port, share):
                     label="Site Type",
                     value="Primary",
                 )
+                sex_dropdown = gr.Dropdown(
+                    choices=SEX_OPTIONS,
+                    label="Sex",
+                    value="Unknown",
+                )
+                tissue_site_dropdown = gr.Dropdown(
+                    choices=get_tissue_sites(),
+                    label="Tissue Site",
+                    value="Unknown",
+                )
                 cancer_subtype_dropdown = gr.Dropdown(
                     choices=[name for name in cancer_subtype_name_map.keys()],
                     label="Cancer Subtype",
@@ -195,15 +209,9 @@ def launch_gradio(server_name, server_port, share):
                 )
                 with gr.Row():
                     settings_input = gr.Dataframe(
-                        headers=[
-                            "Slide",
-                            "Site Type",
-                            "Cancer Subtype",
-                            "IHC Subtype",
-                            "Segmentation Config",
-                        ],
+                        headers=SETTINGS_COLUMNS,
                         label="Current Settings",
-                        datatype=["str", "str", "str", "str", "str"],
+                        datatype=["str"] * len(SETTINGS_COLUMNS),
                         visible=False,
                         interactive=True,
                         static_columns="Slide",
@@ -270,7 +278,7 @@ def launch_gradio(server_name, server_port, share):
                 gr.File(visible=False),
             )
 
-        def get_settings(files, site_type, cancer_subtype, ihc_subtype, seg_config):
+        def get_settings(files, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config):
             if files is None:
                 return pd.DataFrame()
             settings = []
@@ -278,7 +286,7 @@ def launch_gradio(server_name, server_port, share):
                 filename = file.name if hasattr(file, "name") else file
                 slide_name = filename.split("/")[-1]
                 settings.append(
-                    [slide_name, site_type, cancer_subtype, ihc_subtype, seg_config]
+                    [slide_name, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config]
                 )
             df = pd.DataFrame(settings, columns=SETTINGS_COLUMNS)
             return df
@@ -288,6 +296,8 @@ def launch_gradio(server_name, server_port, share):
             [
                 input_slides.change,
                 site_dropdown.change,
+                sex_dropdown.change,
+                tissue_site_dropdown.change,
                 cancer_subtype_dropdown.change,
                 ihc_subtype_dropdown.change,
                 seg_config_dropdown.change,
@@ -295,18 +305,20 @@ def launch_gradio(server_name, server_port, share):
             inputs=[
                 input_slides,
                 site_dropdown,
+                sex_dropdown,
+                tissue_site_dropdown,
                 cancer_subtype_dropdown,
                 ihc_subtype_dropdown,
                 seg_config_dropdown,
             ],
             outputs=[settings_input, settings_csv, ihc_subtype_dropdown],
         )
-        def update_settings(files, site_type, cancer_subtype, ihc_subtype, seg_config):
+        def update_settings(files, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config):
             has_ihc = "Breast" in cancer_subtype
             if not files:
                 return None, None, gr.Dropdown(visible=has_ihc)
             settings_df = get_settings(
-                files, site_type, cancer_subtype, ihc_subtype, seg_config
+                files, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config
             )
             if settings_df is not None:
                 has_ihc = any("Breast" in cs for cs in settings_df["Cancer Subtype"])

src/mosaic/ui/utils.py

@@ -17,6 +17,7 @@ import requests
 TEMP_USER_DATA_DIR = Path(tempfile.gettempdir()) / "mosaic_user_data"
 
 IHC_SUBTYPES = ["", "HR+/HER2+", "HR+/HER2-", "HR-/HER2+", "HR-/HER2-"]
+SEX_OPTIONS = ["Unknown", "Male", "Female"]
 
 SETTINGS_COLUMNS = [
     "Slide",
@@ -29,6 +30,23 @@ SETTINGS_COLUMNS = [
 ]
 
 oncotree_code_map = {}
+tissue_site_list = None
+
+
+def get_tissue_sites():
+    """Get the list of tissue sites from the tissue site map file.
+    
+    Returns:
+        List of tissue site names
+    """
+    global tissue_site_list
+    if tissue_site_list is None:
+        current_dir = Path(__file__).parent.parent.parent
+        tissue_site_map_path = current_dir / "data" / "tissue_site_original_to_idx.csv"
+        df = pd.read_csv(tissue_site_map_path)
+        # Get unique tissue sites and sort them
+        tissue_site_list = ["Unknown"] + sorted(df["TISSUE_SITE"].unique().tolist())
+    return tissue_site_list
 
 
 def get_oncotree_code_name(code):
@@ -100,6 +118,10 @@ def load_settings(slide_csv_path):
         settings_df["Cancer Subtype"] = "Unknown"
     if "IHC Subtype" not in settings_df.columns:
         settings_df["IHC Subtype"] = ""
+    if "Sex" not in settings_df.columns:
+        settings_df["Sex"] = "Unknown"
+    if "Tissue Site" not in settings_df.columns:
+        settings_df["Tissue Site"] = "Unknown"
     if not set(SETTINGS_COLUMNS).issubset(settings_df.columns):
         raise ValueError("Missing required column in CSV file")
     settings_df = settings_df[SETTINGS_COLUMNS]