Update app.py

app.py

@@ -146,6 +146,7 @@ def load_model(path):
 
 # --------- Gradio Functions --------- #
 def train_model(font_file, epochs=100, learning_rate=0.001):
+    import time
     global font_path, ocr_model
 
     # Save the uploaded font file
@@ -159,41 +160,54 @@ def train_model(font_file, epochs=100, learning_rate=0.001):
     dataset = OCRDataset(font_path)
     dataloader = DataLoader(dataset, batch_size=16, shuffle=True, collate_fn=custom_collate_fn)
 
-    # Visualize one sample for sanity check
+    # Visualize one sample
     img, label = dataset[0]
     print("Label:", ''.join([IDX2CHAR[i.item()] for i in label]))
     plt.imshow(img.permute(1, 2, 0).squeeze(), cmap='gray')
     plt.show()
 
-    # Initialize model
+    # Initialize model, loss, optimizer, scheduler
     model = OCRModel(num_classes=len(CHAR2IDX)).to(device)
-    criterion = nn.CTCLoss(blank=0)
-    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+    criterion = nn.CTCLoss(blank=BLANK_IDX)
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)
 
     # Training loop
     for epoch in range(epochs):
+        model.train()
+        running_loss = 0.0
+
         for img, targets, target_lengths in dataloader:
             img = img.to(device)
             targets = targets.to(device)
             target_lengths = target_lengths.to(device)
 
-            output = model(img)
-            batch_size = img.size(0)
+            output = model(img)  # [B, T, C]
             seq_len = output.size(1)
-            input_lengths = torch.full(size=(batch_size,), fill_value=seq_len, dtype=torch.long).to(device)
+            batch_size = img.size(0)
+            input_lengths = torch.full((batch_size,), seq_len, dtype=torch.long).to(device)
+
+            log_probs = output.log_softmax(2).transpose(0, 1)  # [T, B, C]
+            loss = criterion(log_probs, targets, input_lengths, target_lengths)
 
-            loss = criterion(output.log_softmax(2).transpose(0, 1), targets, input_lengths, target_lengths)
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
 
-        print(f"Epoch {epoch + 1}, Loss: {loss.item():.4f}")
+            running_loss += loss.item()
 
-    # Save model with structured name
-    model_name = f"{font_name}_{epochs}epochs_lr{learning_rate:.0e}.pth"
+        avg_loss = running_loss / len(dataloader)
+        scheduler.step(avg_loss)
+        print(f"[{epoch + 1}/{epochs}] Loss: {avg_loss:.4f}")
+
+    # Save the trained model
+    timestamp = time.strftime("%Y%m%d%H%M%S")
+    model_name = f"{font_name}_{epochs}ep_lr{learning_rate:.0e}_{timestamp}.pth"
     save_model(model, model_name)
     ocr_model = model
-    return f"Training complete! Model saved as '{model_name}'."
+
+    return f"✅ Training complete! Model saved as '{model_name}'"
+