Update app.py

app.py

@@ -41,6 +41,7 @@ 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.
@@ -50,7 +51,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)
     
-    # Make the figure itself smaller
+    # Set a smaller figure size to produce a smaller output image.
     plt.figure(figsize=(4, 4))
     sns.heatmap(matrix,
                 mask=mask,
@@ -73,6 +74,8 @@ def generate_heatmap_image(model_entry):
     
     # 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((300, 300), Image.LANCZOS)
     return image
 
 def generate_heatmaps(selected_model_names):
@@ -81,10 +84,7 @@ 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 = []
-    for entry in filtered_entries:
-        img = generate_heatmap_image(entry)
-        images.append(img)
+    images = [generate_heatmap_image(entry) for entry in filtered_entries]
     return images
 
 # -------------------------------
@@ -104,13 +104,7 @@ with gr.Blocks() as demo:
     
     generate_btn = gr.Button("Generate Heatmaps")
     # The 'columns' parameter will display images in a grid with 2 columns.
-    # 'image_size=(300, 300)' ensures each image is displayed at 300x300 px.
-    gallery = gr.Gallery(
-        label="Heatmaps",
-        columns=2,
-        image_size=(200, 200),
-        object_fit="contain"
-    )
+    gallery = gr.Gallery(label="Heatmaps", columns=2)
     
     generate_btn.click(fn=generate_heatmaps, inputs=model_dropdown, outputs=gallery)