Update app.py

app.py

@@ -13,150 +13,236 @@ from sklearn.linear_model import LogisticRegression
 from huggingface_hub import hf_hub_download
 
 
-def load_model_params(model_path):
-    """Load model parameters from a JSON file."""
+def load_and_predict_with_classifier(x, model_path, output_path, save):
+
+    # Load the model parameters from the JSON file
     with open(model_path, 'r') as f:
         model_params = json.load(f)
-    return model_params
-
-def reconstruct_classifier(model_params):
-    """Reconstruct the logistic regression model from parameters."""
-    model = LogisticRegression(multi_class='multinomial', solver='lbfgs', max_iter=1000)
-    model.coef_ = np.array(model_params["coef"])
-    model.intercept_ = np.array(model_params["intercept"])
-    model.classes_ = np.array(model_params["classes"])
-    return model
-
-def save_predictions(y_pred, output_path):
-    """Save predictions to a CSV file."""
-    df = pd.DataFrame(y_pred, columns=["predicted_cell_type"])
-    df.to_csv(output_path, index=False, header=False)
-
-def load_and_predict_with_classifier(x, model_path, output_path, save=False):
-    """Load model, predict, and optionally save predictions."""
-    model_params = load_model_params(model_path)
-    model = reconstruct_classifier(model_params)
-    y_pred = model.predict(x)
+
+    # 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:
-        save_predictions(y_pred, output_path)
+        df = pd.DataFrame(y_pred, columns=["predicted_cell_type"])
+        df.to_csv(output_path, index=False, header=False)
+    
     return y_pred
 
+
 def plot_umap(adata):
-    """Generate a UMAP plot from the provided AnnData object."""
+
     labels = pd.Categorical(adata.obs["cell_type"])
+
     reducer = umap.UMAP(n_neighbors=15, min_dist=0.1, n_components=2, random_state=42)
     embedding = reducer.fit_transform(adata.obsm["X_uce"])
 
     plt.figure(figsize=(10, 8))
+
+    # Create the scatter plot
     scatter = plt.scatter(embedding[:, 0], embedding[:, 1], c=labels.codes, cmap='Set1', s=50, alpha=0.6)
 
-    handles = [
-        plt.Line2D([0], [0], marker='o', color='w', label=cell_type,
-                   markerfacecolor=plt.cm.Set1(i / len(labels.categories)), markersize=10)
-        for i, cell_type in enumerate(labels.categories)
-    ]
+    # Create a legend
+    handles = []
+    for i, cell_type in enumerate(labels.categories):
+        handles.append(plt.Line2D([0], [0], marker='o', color='w', label=cell_type,
+                                  markerfacecolor=plt.cm.Set1(i / len(labels.categories)), markersize=10))
+
     plt.legend(handles=handles, title='Cell Type')
     plt.title('UMAP projection of the data')
     plt.xlabel('UMAP1')
     plt.ylabel('UMAP2')
-
+    
+    # Save plot to a BytesIO object
     buf = BytesIO()
     plt.savefig(buf, format='png')
     buf.seek(0)
+    
+    # Read the image from BytesIO object
     img = plt.imread(buf, format='png')
+
     return img
 
+
 def toggle_file_input(default_dataset):
-    """Toggle file input based on dataset selection."""
     if default_dataset != "None":
-        return gr.update(interactive=False)
+        return gr.update(interactive=False)  # Disable the file input if a default dataset is selected
     else:
-        return gr.update(interactive=True)
+        return gr.update(interactive=True)  # Enable the file input if no default dataset is selected
 
-def run_uce_model(input_file_path, model_dir, model_loc):
-    """Run UCE model on the provided AnnData file."""
-    command = [
-        sys.executable, 
-        os.path.join(model_dir, 'eval_single_anndata.py'), 
-        '--adata_path', input_file_path, 
-        '--dir', model_dir, 
-        '--model_loc', model_loc
-    ]
-    subprocess.run(command, check=True)
 
 def main(input_file_path, species, default_dataset):
-    """Main function to execute the demo logic."""
 
-    # Clone the UCE repository and set paths
-    repo_url = 'https://github.com/minwoosun/UCE.git'
-    repo_dir = '/home/user/app/UCE'
-    if not os.path.exists(repo_dir):
-        subprocess.run(['git', 'clone', repo_url], check=True)
+    # Get the current working directory
+    current_working_directory = os.getcwd()
     
-    sys.path.append(repo_dir)
-    
-    # Handle default datasets
-    default_dataset_paths = {
-        "PBMC 100 cells": hf_hub_download(repo_id="minwoosun/uce-misc", filename="100_pbmcs_proc_subset.h5ad"),
-        "PBMC 1000 cells": hf_hub_download(repo_id="minwoosun/uce-misc", filename="1k_pbmcs_proc_subset.h5ad"),
-    }
+    # Print the current working directory
+    print("Current Working Directory:", current_working_directory)
+
+    # clone and cd into UCE repo
+    os.system('git clone https://github.com/minwoosun/UCE.git')
+    os.chdir('/home/user/app/UCE')
+
+    # Get the current working directory
+    current_working_directory = os.getcwd()
     
-    if default_dataset in default_dataset_paths:
-        input_file_path = default_dataset_paths[default_dataset]
+    # Print the current working directory
+    print("Current Working Directory:", current_working_directory)
 
-    # Run UCE model
-    run_uce_model(input_file_path, repo_dir, 'minwoosun/uce-100m')
+    # Specify the path to the directory you want to add
+    new_directory = "/home/user/app/UCE"
+    
+    # Add the directory to the Python path
+    sys.path.append(new_directory)
+
+    # Set default dataset path
+    default_dataset_1_path = hf_hub_download(repo_id="minwoosun/uce-misc", filename="100_pbmcs_proc_subset.h5ad")
+    default_dataset_2_path = hf_hub_download(repo_id="minwoosun/uce-misc", filename="1k_pbmcs_proc_subset.h5ad")
+
+    # If the user selects a default dataset, use that instead of the uploaded file
+    if default_dataset == "PBMC 100 cells":
+        input_file_path = default_dataset_1_path
+    elif default_dataset == "PBMC 1000 cells":
+        input_file_path = default_dataset_2_path
+
+    ##############
+    #     UCE    #
+    ##############
+    from evaluate import AnndataProcessor
+    from accelerate import Accelerator
+
+    dir_path = '/home/user/app/UCE/'
+    model_loc = 'minwoosun/uce-100m'
+
+    print(input_file_path)
+    print(dir_path)
+    print(model_loc)
+    
+    # Construct the command
+    command = [
+        'python', 
+        '/home/user/app/UCE/eval_single_anndata.py', 
+        '--adata_path', input_file_path, 
+        '--dir', dir_path, 
+        '--model_loc', model_loc
+    ]
+    
+    # Print the command for debugging
+    print("Running command:", command)
+
+    print("---> RUNNING UCE")
+    result = subprocess.run(command, capture_output=True, text=True, check=True)
+    print(result.stdout)
+    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")
 
-    # Load UCE embeddings and perform classification
-    adata = sc.read_h5ad(os.path.join(repo_dir, f"{os.path.splitext(os.path.basename(input_file_path))[0]}_uce_adata.h5ad"))
+    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']
 
-    model_path = hf_hub_download(repo_id="minwoosun/uce-misc", filename="tabula_sapiens_v1_logistic_regression_model_weights.json")
-    pred_file = os.path.join(repo_dir, f"{os.path.splitext(os.path.basename(input_file_path))[0]}_predictions.csv")
     y_pred = load_and_predict_with_classifier(x, model_path, pred_file, save=True)
-
-    # Generate UMAP plot
+    
+    ##############
+    #    UMAP    #
+    ##############
     img = plot_umap(adata)
+  
+    return img, output_file, pred_file
 
-    return img, os.path.join(repo_dir, f"{os.path.splitext(os.path.basename(input_file_path))[0]}_uce_adata.h5ad"), pred_file
 
-# Gradio UI
+if __name__ == "__main__":   
 
-def create_demo():
-    """Create and launch the Gradio demo."""
-    
     with gr.Blocks() as demo:
-        gr.Markdown("""
+        gr.Markdown(
+            '''
             <div style="text-align:center; margin-bottom:20px;">
-                <h1>UCE 100M Demo 🦠</h1>
-                <h2>Universal Cell Embeddings: Zero-Shot Cell-Type Classification in Action!</h2>
-                <div style="margin-top:10px;">
-                    <a href="https://github.com/minwoosun/UCE"><img src="https://badges.aleen42.com/src/github.svg" alt="GitHub"></a>
-                    <a href="https://www.biorxiv.org/content/10.1101/2023.11.28.568918v1"><img src="https://img.shields.io/badge/bioRxiv-2023.11.28.568918-green?style=plastic" alt="Paper"></a>
-                    <a href="https://colab.research.google.com/drive/1opud0BVWr76IM8UnGgTomVggui_xC4p0?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
-                </div>
-                <p>Upload a `.h5ad` single cell gene expression file or select the species to generate UMAP projections and download the embeddings.</p>
+                <span style="font-size:3em; font-weight:bold;">UCE 100M Demo 🦠</span>
             </div>
-        """)
-        
-        # Inputs
+            <div style="text-align:center; margin-bottom:10px;">
+                <span style="font-size:1.5em; font-weight:bold;">Universal Cell Embeddings: Zero-Shot Cell-Type Classification in Action!</span>
+            </div>
+            <div style="text-align:center; margin-bottom:20px;">
+                <a href="https://github.com/minwoosun/UCE">
+                    <img src="https://badges.aleen42.com/src/github.svg" alt="GitHub" style="display:inline-block; margin-right:10px;">
+                </a>
+                <a href="https://www.biorxiv.org/content/10.1101/2023.11.28.568918v1">
+                    <img src="https://img.shields.io/badge/bioRxiv-2023.11.28.568918-green?style=plastic" alt="Paper" style="display:inline-block; margin-right:10px;">
+                </a>
+                <a href="https://colab.research.google.com/drive/1opud0BVWr76IM8UnGgTomVggui_xC4p0?usp=sharing">
+                    <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" style="display:inline-block; margin-right:10px;">
+                </a>
+            </div>
+            <div style="text-align:left; margin-bottom:20px;">
+                Upload a `.h5ad` single cell gene expression file and select the species (Human/Mouse). 
+                The demo will generate UMAP projections of the embeddings and allow you to download the embeddings for further analysis.
+            </div>
+            <div style="margin-bottom:20px;">
+                <ol style="list-style:none; padding-left:0;">
+                    <li>1. Upload your `.h5ad` file or select one of the default datasets (subset of 10x pbmc data)</li>
+                    <li>2. Select the species</li>
+                    <li>3. Click "Run" to view the UMAP scatter plot</li>
+                    <li>4. Download the UCE embeddings and predicted cell-types</li>
+                </ol>
+            </div>
+            <div style="text-align:left; line-height:1.8;">
+                Please consider citing the following paper if you use this tool in your research:
+            </div>
+            <div style="text-align:left; line-height:1.8;">
+                Rosen, Y., Roohani, Y., Agarwal, A., Samotorčan, L., Tabula Sapiens Consortium, Quake, S. R., & Leskovec, J. Universal Cell Embeddings: A Foundation Model for Cell Biology. bioRxiv. https://doi.org/10.1101/2023.11.28.568918
+            </div>
+            '''
+        )
+
+        # Define Gradio inputs and outputs
         file_input = gr.File(label="Upload a .h5ad single cell gene expression file or select a default dataset below")
-        species_input = gr.Dropdown(choices=["human", "mouse"], label="Select species")
-        default_dataset_input = gr.Dropdown(choices=["None", "PBMC 100 cells", "PBMC 1000 cells"], label="Select default dataset")
-        
-        default_dataset_input.change(toggle_file_input, inputs=[default_dataset_input], outputs=[file_input])
+        # species_input = gr.Dropdown(choices=["human", "mouse"], label="Select species")
+        with gr.Row():
+            species_input = gr.Dropdown(choices=["human", "mouse"], label="Select species")
+            default_dataset_input = gr.Dropdown(choices=["None", "PBMC 100 cells", "PBMC 1000 cells"], label="Select default dataset")
+
+        # Attach the `change` event to the dropdown
+        default_dataset_input.change(
+            toggle_file_input,
+            inputs=[default_dataset_input],
+            outputs=[file_input]
+        )
+
+        run_button = gr.Button("Run", elem_classes="run-button")
         
-        # Outputs
-        run_button = gr.Button("Run")
+        # Arrange UMAP plot and file output side by side
         with gr.Row():
-            image_output = gr.Image(type="numpy", label="UMAP of UCE Embeddings")
+            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")
+
+            print(image_output)
+            print(file_output)
+            print(pred_output)
     
-        # Run the function on button click
-        run_button.click(fn=main, inputs=[file_input, species_input, default_dataset_input], outputs=[image_output, file_output, pred_output])
+        # Add the components and link to the function
+        run_button.click(
+            fn=main, 
+            inputs=[file_input, species_input, default_dataset_input], 
+            outputs=[image_output, file_output, pred_output]
+        )
 
-    demo.launch()
 
-if __name__ == "__main__":
-    create_demo()
+    demo.launch()