Incorporate classifier

app.py

@@ -1,17 +1,45 @@
 import gradio as gr
 import pandas as pd
+import numpy as np
 import umap
+import json
 import matplotlib.pyplot as plt
 import os
-import tempfile
+# import tempfile
 import scanpy as sc
-import argparse
+# import argparse
 import subprocess
 import sys
+from evaluate import AnndataProcessor
+from accelerate import Accelerator
 from io import BytesIO
+from sklearn.linear_model import LogisticRegression
 from huggingface_hub import hf_hub_download
 
 
+def load_and_predict_with_classifier(x, model_path, output_path):
+
+    # Load the model parameters from the JSON file
+    with open(model_path, 'r') as f:
+        model_params = json.load(f)
+
+    # Reconstruct the logistic regression model
+    model_loaded = LogisticRegression(multi_class='multinomial', solver='lbfgs', max_iter=1000)
+    model_loaded.coef_ = np.array(model_params["coef"])
+    model_loaded.intercept_ = np.array(model_params["intercept"])
+    model_loaded.classes_ = np.array(model_params["classes"])
+
+    # output predictions
+    y_pred = model_loaded.predict(x)
+
+    # Convert the array to a Pandas DataFrame
+    if save:
+        df = pd.DataFrame(y_pred, columns=["predicted_cell_type"])
+        df.to_csv(output_path, index=False, header=False)
+    
+    return y_pred
+
+
 def main(input_file_path, species):
 
     # Get the current working directory
@@ -40,8 +68,6 @@ def main(input_file_path, species):
     ##############
     #     UCE    #
     ##############
-    from evaluate import AnndataProcessor
-    from accelerate import Accelerator
 
     # #  python eval_single_anndata.py --adata_path "./data/10k_pbmcs_proc.h5ad" --dir "./" --model_loc "minwoosun/uce-100m"
     # script_name = "/home/user/app/UCE/eval_single_anndata.py"
@@ -77,6 +103,25 @@ def main(input_file_path, species):
     print(result.stderr)
     print("---> FINSIH UCE")
 
+    ################################
+    #   Cell-type classification   #
+    ################################
+
+    
+    # Set output file path
+    file_name_with_ext = os.path.basename(input_file_path)
+    file_name = os.path.splitext(file_name_with_ext)[0]
+    pred_file = "/home/user/app/UCE/" + f"{file_name}_predictions.csv"
+    model_path = hf_hub_download(repo_id="minwoosun/uce-misc", filename="tabula_sapiens_v1_logistic_regression_model_weights.json")
+
+    file_name_with_ext = os.path.basename(input_file_path)
+    file_name = os.path.splitext(file_name_with_ext)[0]
+    output_file = "/home/user/app/UCE/" + f"{file_name}_uce_adata.h5ad"
+    adata = sc.read_h5ad(output_file)
+    x = adata.obsm['X_uce']
+
+    y_pred = load_and_predict_with_classifier(x, model_path, pred_file, save=True)
+    
     ##############
     #    UMAP    #
     ##############
@@ -84,12 +129,12 @@ def main(input_file_path, species):
 
     if (UMAP):
 
-        # Set output file path
-        file_name_with_ext = os.path.basename(input_file_path)
-        file_name = os.path.splitext(file_name_with_ext)[0]
-        output_file = "/home/user/app/UCE/" + f"{file_name}_uce_adata.h5ad"
+        # # Set output file path
+        # file_name_with_ext = os.path.basename(input_file_path)
+        # file_name = os.path.splitext(file_name_with_ext)[0]
+        # output_file = "/home/user/app/UCE/" + f"{file_name}_uce_adata.h5ad"
         
-        adata = sc.read_h5ad(output_file)
+        # adata = sc.read_h5ad(output_file)
 
         labels = pd.Categorical(adata.obs["cell_type"])
 
@@ -123,7 +168,7 @@ def main(input_file_path, species):
         img = None
         print("no image")
     
-    return img, output_file
+    return img, output_file, pred_file
 
 
 if __name__ == "__main__":   
@@ -181,12 +226,13 @@ if __name__ == "__main__":
         with gr.Row():
             image_output = gr.Image(type="numpy", label="UMAP of UCE Embeddings")
             file_output = gr.File(label="Download embeddings")
+            pred_output = gr.File(label="Download predictions")
     
         # Add the components and link to the function
         run_button.click(
             fn=main, 
             inputs=[file_input, species_input], 
-            outputs=[image_output, file_output]
+            outputs=[image_output, file_output, pred_output]
         )
 
     demo.launch()