improved prediction viz

app.py

@@ -290,7 +290,7 @@ def run_inference(prompt_cell, prompt_cond, query_cell):
     
     time.sleep(0.5)
     
-    # Result state - show predicted cells with the new condition applied
+    # Result state - show predicted gene counts per query cell
     predicted_cells = generate_cell_array(query_cell, prompt_cond, 5, is_prompt=True)
     
     result_html = f'''
@@ -300,43 +300,58 @@ def run_inference(prompt_cell, prompt_cond, query_cell):
         <div style="text-align: center; margin-bottom: 24px;">
             <div style="font-family: monospace; font-size: 10px; color: #64748b; letter-spacing: 0.15em; margin-bottom: 8px; text-transform: uppercase;">PREDICTION COMPLETE</div>
             <div style="font-size: 14px; font-weight: 600; color: #1e293b;">
-                {query_cell_data["name"]} gene expression under {prompt_cond_data["name"]}
+                Predicted gene expression counts for {query_cell_data["name"]}
+            </div>
+            <div style="font-size: 11px; color: #64748b; margin-top: 4px;">
+                under <span style="color: {prompt_cond_data["color"]}; font-weight: 600;">{prompt_cond_data["name"]}</span> condition
             </div>
         </div>
         
-        <!-- Predicted Cells -->
-        <div style="display: flex; justify-content: center; margin-bottom: 24px;">
-            <div style="background: {prompt_cond_data["bg"]}20; border: 2px solid {prompt_cond_data["color"]}60; border-radius: 12px; padding: 16px 24px; display: flex; gap: 8px; align-items: center;">
-                {predicted_cells}
-            </div>
+        <!-- Section label -->
+        <div style="display: flex; align-items: center; gap: 12px; margin-bottom: 16px; justify-content: center;">
+            <div style="height: 1px; width: 60px; background: linear-gradient(90deg, transparent, #cbd5e1);"></div>
+            <div style="font-size: 10px; font-weight: 600; color: #475569; text-transform: uppercase; letter-spacing: 0.08em;">Predicted Gene Counts Per Query Cell</div>
+            <div style="height: 1px; width: 60px; background: linear-gradient(90deg, #cbd5e1, transparent);"></div>
         </div>
         
-        <!-- Output columns visualization -->
-        <div style="display: flex; justify-content: center; gap: 12px; margin-bottom: 20px;">
+        <!-- Output columns visualization - Gene counts per cell -->
+        <div style="display: flex; justify-content: center; gap: 16px; margin-bottom: 20px;">
             <div style="text-align: center;">
-                <div style="background: rgba(99, 102, 241, 0.2); border: 1px solid rgba(99, 102, 241, 0.4); border-radius: 8px; padding: 8px 12px; display: flex; flex-direction: column; gap: 4px;">
+                <div style="background: white; border: 2px solid {query_cell_data["color"]}40; border-radius: 10px; padding: 10px 14px; display: flex; flex-direction: column; gap: 2px; box-shadow: 0 2px 8px rgba(0,0,0,0.05);">
                     {generate_output_column(prompt_cond_data["color"])}
                 </div>
-                <div style="font-size: 10px; color: #64748b; margin-top: 4px;">c₁</div>
+                <div style="display: flex; align-items: center; justify-content: center; gap: 4px; margin-top: 6px;">
+                    {make_cell_svg(query_cell_data["color"], False, 18)}
+                    <div style="font-size: 9px; color: #475569; font-weight: 500;">Cell 1</div>
+                </div>
             </div>
             <div style="text-align: center;">
-                <div style="background: rgba(99, 102, 241, 0.2); border: 1px solid rgba(99, 102, 241, 0.4); border-radius: 8px; padding: 8px 12px; display: flex; flex-direction: column; gap: 4px;">
+                <div style="background: white; border: 2px solid {query_cell_data["color"]}40; border-radius: 10px; padding: 10px 14px; display: flex; flex-direction: column; gap: 2px; box-shadow: 0 2px 8px rgba(0,0,0,0.05);">
                     {generate_output_column(prompt_cond_data["color"])}
                 </div>
-                <div style="font-size: 10px; color: #64748b; margin-top: 4px;">c₂</div>
+                <div style="display: flex; align-items: center; justify-content: center; gap: 4px; margin-top: 6px;">
+                    {make_cell_svg(query_cell_data["color"], False, 18)}
+                    <div style="font-size: 9px; color: #475569; font-weight: 500;">Cell 2</div>
+                </div>
             </div>
-            <div style="display: flex; align-items: center; color: #64748b; font-size: 14px; letter-spacing: 3px;">···</div>
+            <div style="display: flex; align-items: center; color: #94a3b8; font-size: 16px; font-weight: bold; letter-spacing: 3px; padding-bottom: 24px;">···</div>
             <div style="text-align: center;">
-                <div style="background: rgba(99, 102, 241, 0.2); border: 1px solid rgba(99, 102, 241, 0.4); border-radius: 8px; padding: 8px 12px; display: flex; flex-direction: column; gap: 4px;">
+                <div style="background: white; border: 2px solid {query_cell_data["color"]}40; border-radius: 10px; padding: 10px 14px; display: flex; flex-direction: column; gap: 2px; box-shadow: 0 2px 8px rgba(0,0,0,0.05);">
                     {generate_output_column(prompt_cond_data["color"])}
                 </div>
-                <div style="font-size: 10px; color: #64748b; margin-top: 4px;">cₙ</div>
+                <div style="display: flex; align-items: center; justify-content: center; gap: 4px; margin-top: 6px;">
+                    {make_cell_svg(query_cell_data["color"], False, 18)}
+                    <div style="font-size: 9px; color: #475569; font-weight: 500;">Cell n</div>
+                </div>
             </div>
         </div>
         
         <!-- Description -->
-        <div style="text-align: center; font-size: 11px; color: #64748b; max-width: 300px; margin: 0 auto; line-height: 1.5;">
-            Zero-shot prediction of gene expression counts using in-context learning from {prompt_cell_data["name"]} response to {prompt_cond_data["name"]}.
+        <div style="text-align: center; font-size: 11px; color: #64748b; max-width: 360px; margin: 0 auto; line-height: 1.6;">
+            <strong style="color: #475569;">Zero-shot prediction:</strong> Using in-context learning from 
+            <span style="color: {prompt_cond_data["color"]}; font-weight: 500;">{prompt_cell_data["name"]}</span> response, 
+            STACK predicts gene counts for each <span style="color: {query_cell_data["color"]}; font-weight: 500;">{query_cell_data["name"]}</span> 
+            under the same perturbation.
         </div>
         
     </div>
@@ -345,11 +360,23 @@ def run_inference(prompt_cell, prompt_cond, query_cell):
     yield result_html, gr.update(visible=True), gr.update(visible=False)
 
 def generate_output_column(color):
-    """Generate a vertical column of gene expression values"""
+    """Generate a vertical column of gene expression counts showing explicit values"""
+    gene_names = ["g₁", "g₂", "g₃", "g₄", "g₅"]
     cells = []
-    for _ in range(5):
-        opacity = 0.2 + random.random() * 0.6
-        cells.append(f'<div style="width: 16px; height: 16px; background: {color}; opacity: {opacity:.1f}; border-radius: 3px;"></div>')
+    for i, gene in enumerate(gene_names):
+        # Generate pseudo gene count value
+        count = random.randint(10, 500)
+        bar_width = min(count / 500 * 40, 40)  # Scale to max 40px
+        opacity = 0.4 + (count / 500) * 0.5
+        cells.append(f'''
+        <div style="display: flex; align-items: center; gap: 4px; height: 18px;">
+            <div style="font-size: 8px; color: #64748b; width: 14px; text-align: right;">{gene}</div>
+            <div style="width: 44px; height: 12px; background: #e2e8f0; border-radius: 2px; overflow: hidden; position: relative;">
+                <div style="width: {bar_width}px; height: 100%; background: {color}; opacity: {opacity:.1f}; border-radius: 2px;"></div>
+            </div>
+            <div style="font-size: 8px; color: #475569; font-family: monospace; width: 22px;">{count}</div>
+        </div>
+        ''')
     return '\n'.join(cells)
 
 def reset_inference(prompt_cell, prompt_cond, query_cell):