create umap function

app.py

@@ -5,13 +5,9 @@ import umap
 import json
 import matplotlib.pyplot as plt
 import os
-# import tempfile
 import scanpy as sc
-# 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
@@ -40,6 +36,40 @@ def load_and_predict_with_classifier(x, model_path, output_path, save):
     return y_pred
 
 
+def plot_umap(adata):
+
+    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))
+
+    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 main(input_file_path, species, default_dataset):
 
     # Get the current working directory
@@ -80,21 +110,12 @@ def main(input_file_path, species, default_dataset):
     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"
-    # args = ["--adata_path", input_file_path, "--dir", "/home/user/app/UCE/", "--model_loc", "minwoosun/uce-100m"]
-    # command = ["python", script_name] + args
-
     dir_path = '/home/user/app/UCE/'
     model_loc = 'minwoosun/uce-100m'
 
     print(input_file_path)
     print(dir_path)
     print(model_loc)
-
-    # # Verify adata_path is not None
-    # if input_file_path is None or not os.path.exists(input_file_path):
-    #     raise ValueError(f"Invalid adata_path: {input_file_path}. Please check if the file exists.")
     
     # Construct the command
     command = [
@@ -118,7 +139,6 @@ def main(input_file_path, species, default_dataset):
     #   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]
@@ -136,49 +156,8 @@ def main(input_file_path, species, default_dataset):
     ##############
     #    UMAP    #
     ##############
-    UMAP = True
-
-    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"
-        
-        # adata = sc.read_h5ad(output_file)
-
-        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))
-
-        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')
-    else:
-        img = None
-        print("no image")
-    
+    img = plot_umap(adata)
+  
     return img, output_file, pred_file