Modularize code

app.py

@@ -13,248 +13,150 @@ from sklearn.linear_model import LogisticRegression
 from huggingface_hub import hf_hub_download
 
 
-def load_and_predict_with_classifier(x, model_path, output_path, save):
-
-    # Load the model parameters from the JSON file
+def load_model_params(model_path):
+    """Load model parameters from a 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
+    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)
     if save:
-        df = pd.DataFrame(y_pred, columns=["predicted_cell_type"])
-        df.to_csv(output_path, index=False, header=False)
-    
+        save_predictions(y_pred, output_path)
     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)
 
-    # 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))
-
+    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)
+    ]
     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)  # Disable the file input if a default dataset is selected
+        return gr.update(interactive=False)
     else:
-        return gr.update(interactive=True)  # Enable the file input if no default dataset is selected
-
-
-def main(input_file_path, species, default_dataset):
-
-    # Get the current working directory
-    current_working_directory = os.getcwd()
-    
-    # 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()
-    
-    # Print the current working directory
-    print("Current Working Directory:", current_working_directory)
-
-    # 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
+        return gr.update(interactive=True)
 
-    dir_path = '/home/user/app/UCE/'
-    model_loc = 'minwoosun/uce-100m'
-
-    print(input_file_path)
-    print(dir_path)
-    print(model_loc)
-    
-    # Construct the command
+def run_uce_model(input_file_path, model_dir, model_loc):
+    """Run UCE model on the provided AnnData file."""
     command = [
-        'python', 
-        '/home/user/app/UCE/eval_single_anndata.py', 
+        sys.executable, 
+        os.path.join(model_dir, 'eval_single_anndata.py'), 
         '--adata_path', input_file_path, 
-        '--dir', dir_path, 
+        '--dir', model_dir, 
         '--model_loc', model_loc
     ]
-    
-    # Print the command for debugging
-    print("Running command:", command)
+    subprocess.run(command, check=True)
 
-    print("---> RUNNING UCE")
-    result = subprocess.run(command, capture_output=True, text=True, check=True)
-    print(result.stdout)
-    print(result.stderr)
-    print("---> FINSIH UCE")
+def main(input_file_path, species, default_dataset):
+    """Main function to execute the demo logic."""
 
-    ################################
-    #   Cell-type classification   #
-    ################################
+    # 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)
+    
+    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"),
+    }
+    
+    if default_dataset in default_dataset_paths:
+        input_file_path = default_dataset_paths[default_dataset]
 
-    # 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")
+    # Run UCE model
+    run_uce_model(input_file_path, repo_dir, 'minwoosun/uce-100m')
 
-    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)
+    # 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"))
     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)
-    
-    ##############
-    #    UMAP    #
-    ##############
+
+    # Generate UMAP plot
     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
 
-if __name__ == "__main__":   
+# Gradio UI
 
+def create_demo():
+    """Create and launch the Gradio demo."""
+    
     with gr.Blocks() as demo:
-        gr.Markdown(
-            '''
-            <div style="text-align:center; margin-bottom:20px;">
-                <span style="font-size:3em; font-weight:bold;">UCE 100M Demo 🦠</span>
-            </div>
-            <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>
+        gr.Markdown("""
             <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:
+                <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>
             </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>
-            '''
-        )
-
-        # download default datasets and assign paths
-        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")
-
-        # Define Gradio inputs and outputs
+        """)
+        
+        # Inputs
         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")
-        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")
+        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")
         
-        # Arrange UMAP plot and file output side by side
+        default_dataset_input.change(toggle_file_input, inputs=[default_dataset_input], outputs=[file_input])
+        
+        # Outputs
+        run_button = gr.Button("Run")
         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)
     
-        # 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]
-        )
-
-        examples = gr.Examples(
-        	examples=[[default_dataset_1_path, "human", "PBMC 100 cells"],[default_dataset_2_path, "human", "PBMC 1000 cells"]],
-        	inputs=[file_input, species_input, default_dataset_input],
-            outputs=[image_output, file_output, pred_output],
-        	fn=main,
-        	cache_examples=True
-        )
+        # 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])
 
     demo.launch()
 
+if __name__ == "__main__":
+    create_demo()