Upload app.py

app.py

@@ -195,21 +195,13 @@ def visualize_predictions(image: Image.Image, predictions: Dict, threshold: floa
     for tag, prob in predictions.get("meta", []):
         if not any(pattern in tag.lower() for pattern in excluded_meta_patterns):
             filtered_meta.append((tag, prob))
-    predictions["meta"] = filtered_meta # Use filtered list for visualization
+    predictions["meta"] = filtered_meta  # Use filtered list for visualization
 
     # --- Plotting Setup ---
-    plt.rcParams['font.family'] = 'DejaVu Sans' # Ensure font compatibility
-    fig = plt.figure(figsize=(12, 20), dpi=100)
-    gs = fig.add_gridspec(1, 2, width_ratios=[1.2, 1])
-
-    # Left side: Image
-    # ax_img = fig.add_subplot(gs[0, 0])
-    # ax_img.imshow(image)
-    # ax_img.set_title("Original Image")
-    # ax_img.axis('off')
-
-    # Right side: Tags
-    ax_tags = fig.add_subplot(gs[0, 1])
+    plt.rcParams['font.family'] = 'DejaVu Sans'
+    fig = plt.figure(figsize=(8, 20), dpi=100)
+    ax_tags = fig.add_subplot(1, 1, 1)
+
     all_tags, all_probs, all_colors = [], [], []
     color_map = {
         'rating': 'red', 'character': 'blue', 'copyright': 'purple',
@@ -223,22 +215,18 @@ def visualize_predictions(image: Image.Image, predictions: Dict, threshold: floa
         ('artist', 'A', color_map['artist']), ('general', 'G', color_map['general']),
         ('meta', 'M', color_map['meta'])
     ]:
-        # Sort within category by probability before adding
         sorted_tags = sorted(predictions.get(cat, []), key=lambda x: x[1], reverse=True)
         for tag, prob in sorted_tags:
-            # Add prefix to tag name for display
-            all_tags.append(f"[{prefix}] {tag.replace('_', ' ')}") # Replace underscores for display
+            all_tags.append(f"[{prefix}] {tag.replace('_', ' ')}")
             all_probs.append(prob)
             all_colors.append(color)
 
     if not all_tags:
         ax_tags.text(0.5, 0.5, "No tags found above threshold", ha='center', va='center')
-        ax_tags.set_title(f"Tags (Thresholds: Gen/Meta={threshold:.2f}, Char/Art/Copy={threshold:.2f})") # Assuming same threshold for now
+        ax_tags.set_title(f"Tags (Threshold ≳ {threshold:.2f})")
         ax_tags.axis('off')
     else:
-        # Sort all aggregated tags by probability (descending) for plotting order
-        # Plotting from bottom up, so we want highest probability at the top
-        sorted_indices = sorted(range(len(all_probs)), key=lambda i: all_probs[i]) # Sort ascending for barh
+        sorted_indices = sorted(range(len(all_probs)), key=lambda i: all_probs[i])
         all_tags = [all_tags[i] for i in sorted_indices]
         all_probs = [all_probs[i] for i in sorted_indices]
         all_colors = [all_colors[i] for i in sorted_indices]
@@ -252,37 +240,31 @@ def visualize_predictions(image: Image.Image, predictions: Dict, threshold: floa
         ax_tags.set_yticklabels(all_tags)
 
         fontsize = 10 if num_tags <= 40 else 8 if num_tags <= 60 else 6
-        for label in ax_tags.get_yticklabels():
-            label.set_fontsize(fontsize)
+        for lbl in ax_tags.get_yticklabels():
+            lbl.set_fontsize(fontsize)
 
-        # Add probability text next to bars
         for i, (bar, prob) in enumerate(zip(bars, all_probs)):
-             # Position text slightly outside the bar, ensuring it stays within plot bounds
-             text_x = min(prob + 0.02, 0.98) # Adjust x position
-             ax_tags.text(text_x, y_positions[i], f"{prob:.3f}", va='center', fontsize=fontsize)
+            text_x = min(prob + 0.02, 0.98)
+            ax_tags.text(text_x, y_positions[i], f"{prob:.3f}", va='center', fontsize=fontsize)
 
         ax_tags.set_xlim(0, 1)
-        ax_tags.set_title(f"Tags (Thresholds approx: {threshold:.2f})") # Indicate threshold used
+        ax_tags.set_title(f"Tags (Threshold ≳ {threshold:.2f})")
 
-        # Add legend
         from matplotlib.patches import Patch
         legend_elements = [
-             Patch(facecolor=color, label=cat.capitalize()) for cat, color in color_map.items()
-             if any(t.startswith(f"[{cat[0].upper() if cat != 'copyright' else '©'}]") for t in all_tags)
+            Patch(facecolor=color, label=cat.capitalize())
+            for cat, color in color_map.items()
+            if any(t.startswith(f"[{cat[0].upper() if cat!='copyright' else '©'}]") for t in all_tags)
         ]
         if legend_elements:
-             ax_tags.legend(handles=legend_elements, loc='lower right', fontsize=8)
+            ax_tags.legend(handles=legend_elements, loc='lower right', fontsize=8)
 
     plt.tight_layout()
-    plt.subplots_adjust(bottom=0.05)
-
-    # Save plot to buffer
     buf = io.BytesIO()
     plt.savefig(buf, format='png', dpi=100)
     plt.close(fig)
     buf.seek(0)
-    viz_image = Image.open(buf)
-    return viz_image
+    return Image.open(buf)
 
 # --- Constants ---
 REPO_ID = "cella110n/cl_tagger"