Update app.py

app.py

@@ -240,10 +240,46 @@ def update_class_counts_handler(class_df, dataset_info):
     for _, row in class_df.iterrows():
         if not row["Remove"]:
             rename_to = row["Rename To"]
-            # Aggregate counts for classes that are merged
-            merged_summary[rename_to] = merged_summary.get(rename_to, 0) + updated_counts.get(rename_to, 0)
+            # This logic needs to be careful: sum counts of all original classes that map to the same `rename_to`
+            # Let's recalculate based on mapping
+            merged_summary[rename_to] = 0 # reset
     
-    summary_df = pd.DataFrame(list(merged_summary.items()), columns=["Final Class Name", "Est. Total Images"])
+    for original_name, rename_to in class_mapping.items():
+         if rename_to in merged_summary:
+            # find count for original name in its original mapped state
+            original_count = gather_class_counts(dataset_info, {k:k for k in class_mapping.keys()}).get(original_name,0)
+            is_removed = class_df.loc[class_df['Original Name'] == original_name, 'Remove'].iloc[0]
+            if not is_removed:
+                merged_summary[rename_to] += original_count
+
+    final_summary = {}
+    # Recalculate from scratch for simplicity and accuracy
+    class_map_for_summary = dict(zip(class_df["Original Name"], class_df["Rename To"]))
+    all_final_names = set(class_df[~class_df['Remove']]['Rename To'])
+    
+    final_counts = {name: 0 for name in all_final_names}
+
+    for loc, names, splits, _ in dataset_info:
+        for split in splits:
+            labels_dir = os.path.join(loc, split, 'labels')
+            if not os.path.exists(labels_dir): continue
+            for label_file in os.listdir(labels_dir):
+                found_in_file = set()
+                with open(os.path.join(labels_dir, label_file), 'r') as f:
+                    for line in f:
+                        try:
+                            class_id = int(line.split()[0])
+                            original_name = names[class_id]
+                            is_removed = class_df.loc[class_df['Original Name'] == original_name, 'Remove'].iloc[0]
+                            if not is_removed:
+                                mapped_name = class_map_for_summary.get(original_name)
+                                if mapped_name:
+                                    found_in_file.add(mapped_name)
+                        except (ValueError, IndexError, KeyError): continue
+                for cls in found_in_file:
+                    final_counts[cls] += 1
+    
+    summary_df = pd.DataFrame(list(final_counts.items()), columns=["Final Class Name", "Est. Total Images"])
     return summary_df
 
 def finalize_handler(dataset_info, class_df, progress=gr.Progress()):
@@ -256,7 +292,6 @@ def finalize_handler(dataset_info, class_df, progress=gr.Progress()):
     for _, row in class_df.iterrows():
         if not row["Remove"]:
             rename_to = row["Rename To"]
-            # The limit for a merged class is the sum of the limits of its constituents
             class_limits[rename_to] = class_limits.get(rename_to, 0) + int(row["Max Images"])
             
     status, path = finalize_merged_dataset(dataset_info, class_mapping, class_limits, progress)
@@ -415,7 +450,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
                     epochs_sl = gr.Slider(1, 500, 100, step=1, label="Epochs")
                     batch_sl = gr.Slider(1, 32, 8, step=1, label="Batch Size")
                     imgsz_num = gr.Number(label="Image Size", value=640)
-                    lr_num = gr.Number(label="Learning Rate", value=0.001, format="%.5f")
+                    # <<< FIXED: Removed the 'format' argument which is not supported.
+                    lr_num = gr.Number(label="Learning Rate", value=0.001)
                     opt_dd = gr.Dropdown(["Adam", "AdamW", "SGD"], value="Adam", label="Optimizer")
                     train_btn = gr.Button("Start Training", variant="primary")
                 with gr.Column(scale=2):