Update app.py

app.py

@@ -41,7 +41,6 @@ def generate_heatmap_image(model_entry):
     """
     For a given model entry, extract the six metrics and compute a 6x6 similarity matrix
     using the definition: similarity = 1 - |v_i - v_j|, then return the heatmap as a PIL image.
-    The image is resized to 300x300 pixels.
     """
     scores = model_entry["claude-3.5-sonnet Scores"]["3C3H Scores"]
     # Create a vector with the metrics in the defined order.
@@ -51,8 +50,7 @@ def generate_heatmap_image(model_entry):
     # Create a mask for the upper triangle (keeping the diagonal visible).
     mask = np.triu(np.ones_like(matrix, dtype=bool), k=1)
     
-    # Set a smaller figure size to produce a smaller output image.
-    plt.figure(figsize=(4, 4))
+    plt.figure(figsize=(6, 5))
     sns.heatmap(matrix,
                 mask=mask,
                 annot=True,
@@ -71,11 +69,8 @@ def generate_heatmap_image(model_entry):
     plt.savefig(buf, format="png")
     plt.close()
     buf.seek(0)
-    
-    # Convert the buffer into a PIL Image
+    # Convert the buffer into a PIL Image.
     image = Image.open(buf).convert("RGB")
-    # Resize the image to a fixed size of 300x300 pixels.
-    image = image.resize((600, 600), Image.LANCZOS)
     return image
 
 def generate_heatmaps(selected_model_names):
@@ -84,7 +79,10 @@ def generate_heatmaps(selected_model_names):
     generate a heatmap for each, and return a list of PIL images.
     """
     filtered_entries = [entry for entry in DATA if entry["Meta"]["Model Name"] in selected_model_names]
-    images = [generate_heatmap_image(entry) for entry in filtered_entries]
+    images = []
+    for entry in filtered_entries:
+        img = generate_heatmap_image(entry)
+        images.append(img)
     return images
 
 # -------------------------------
@@ -95,15 +93,10 @@ with gr.Blocks() as demo:
     gr.Markdown("Select the models you want to compare and generate their heatmaps below.")
     
     with gr.Row():
-        model_dropdown = gr.Dropdown(
-            choices=MODEL_NAMES,
-            label="Select Model(s)",
-            multiselect=True,
-            value=MODEL_NAMES[:3]
-        )
+        model_dropdown = gr.Dropdown(choices=MODEL_NAMES, label="Select Model(s)", multiselect=True, value=MODEL_NAMES[:3])
     
     generate_btn = gr.Button("Generate Heatmaps")
-    # The 'columns' parameter will display images in a grid with 2 columns.
+    # Use the 'columns' parameter to set a grid layout in the gallery.
     gallery = gr.Gallery(label="Heatmaps", columns=2)
     
     generate_btn.click(fn=generate_heatmaps, inputs=model_dropdown, outputs=gallery)