Update app.py

app.py

@@ -50,7 +50,6 @@ def aggregate_datasets(api_key, entries):
         imgs, cls_map = fetch_metadata(api_key, ws, proj, ver)
         total_images += imgs
         for cls, cnt in cls_map.items():
-            # normalize class names to lowercase
             norm = cls.strip().lower()
             class_counts[norm] = class_counts.get(norm, 0) + cnt
             class_sources.setdefault(norm, set()).add(url)
@@ -86,7 +85,7 @@ def load_datasets(api_key, file_objs):
         raw = None
         try:
             raw = fobj.read()
-        except:
+        except Exception:
             raw = fobj.get('data') if isinstance(fobj, dict) else None
         if raw is None and isinstance(fobj, str):
             with open(fobj, 'rb') as f:
@@ -118,7 +117,7 @@ def load_datasets(api_key, file_objs):
 def update_classes(df_data):
     """
     Combine edited classes (merge duplicates, lowercase) and recalc.
-    Returns: total_images, plot_fig, json_counts, markdown_summary.
+    Returns: total_images, updated_dataframe, plot_fig, json_counts, markdown_summary.
     """
     combined = {}
     for row in df_data:
@@ -127,14 +126,17 @@ def update_classes(df_data):
         name = row[0].strip().lower()
         try:
             cnt = int(row[1])
-        except:
-            continue
+        except Exception:
+            cnt = 0
         combined[name] = combined.get(name, 0) + cnt
 
     total = sum(combined.values())
+    # Build updated DataFrame
+    updated_df = [[cls, combined[cls]] for cls in combined]
+
     fig = make_bar_chart(combined)
-    md_lines = [f"- **{cls}** ({combined[cls]} images)" for cls in combined]
-    return str(total), fig, json.dumps(combined, indent=2), "\n".join(md_lines)
+    md_summary = "\n".join(f"- **{cls}** ({combined[cls]} images)" for cls in combined)
+    return str(total), updated_df, fig, json.dumps(combined, indent=2), md_summary
 
 
 def build_ui():
@@ -145,21 +147,24 @@ def build_ui():
             api_input = gr.Textbox(label="Roboflow API Key", type="password")
             files = gr.Files(label="Upload .txt files of Roboflow URLs", file_types=[".txt"])
 
-        load = gr.Button("Load Datasets")
+        load_btn = gr.Button("Load Datasets")
         total_out = gr.Textbox(label="Total Images", interactive=False)
-        df = gr.Dataframe(headers=["Class Name", "Count"], interactive=True)
+        df = gr.Dataframe(headers=["Class Name", "Count"], row_count=(1, None), col_count=2, interactive=True)
         plot = gr.Plot()
         json_out = gr.Textbox(label="Counts (JSON)", interactive=False)
         md_out = gr.Markdown(label="Class Sources")
 
-        update = gr.Button("Apply Class Edits")
+        update_btn = gr.Button("Apply Class Edits")
+
+        # Load datasets
+        load_btn.click(fn=load_datasets,
+                       inputs=[api_input, files],
+                       outputs=[total_out, df, plot, json_out, md_out])
 
-        load.click(fn=load_datasets,
-                   inputs=[api_input, files],
-                   outputs=[total_out, df, plot, json_out, md_out])
-        update.click(fn=update_classes,
-                     inputs=[df],
-                     outputs=[total_out, plot, json_out, md_out])
+        # Apply edits and refresh all outputs (including table)
+        update_btn.click(fn=update_classes,
+                         inputs=[df],
+                         outputs=[total_out, df, plot, json_out, md_out])
 
     return demo