Update app.py

app.py

@@ -38,6 +38,10 @@ alphabet = Alphabet.build_alphabet(labels)
 
 # Now initialize decoder correctly
 decoder = BeamSearchDecoderCTC(alphabet)
+# Ensure required directories exist at startup
+os.makedirs("./fonts", exist_ok=True)
+os.makedirs("./models", exist_ok=True)
+os.makedirs("./labels", exist_ok=True)
 
 # --------- Dataset --------- #
 class OCRDataset(Dataset):
@@ -155,7 +159,11 @@ def custom_collate_fn(batch):
 
 # --------- Model Save/Load --------- #
 def list_saved_models():
-    return [f for f in os.listdir() if f.endswith(".pth")]
+    model_dir = "./models"
+    if not os.path.exists(model_dir):
+        return []
+    return [f for f in os.listdir(model_dir) if f.endswith(".pth")]
+
 
 
 def save_model(model, path):
@@ -177,14 +185,18 @@ def train_model(font_file, epochs=100, learning_rate=0.001):
     import time
     global font_path, ocr_model
 
-    # Save uploaded font
+    # Ensure directories exist
+    os.makedirs("./fonts", exist_ok=True)
+    os.makedirs("./models", exist_ok=True)
+
+    # Save uploaded font to ./fonts
     font_name = os.path.splitext(os.path.basename(font_file.name))[0]
-    font_path = f"./{font_name}.ttf"
+    font_path = f"./fonts/{font_name}.ttf"
     with open(font_file.name, "rb") as uploaded:
         with open(font_path, "wb") as f:
             f.write(uploaded.read())
 
-    # Curriculum learning: Start with shorter labels, increase over time
+    # Curriculum learning: label length grows over time
     def get_dataset_for_epoch(epoch):
         if epoch < epochs // 3:
             label_len = (3, 4)
@@ -194,29 +206,27 @@ def train_model(font_file, epochs=100, learning_rate=0.001):
             label_len = (5, 7)
         return OCRDataset(font_path, label_length_range=label_len)
 
-    # Visualize one sample from initial dataset
+    # Visualize one sample
     dataset = get_dataset_for_epoch(0)
-    img, label, _ = dataset[0]  # Ignore the 3rd value (e.g., label length)
-
+    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()
 
-    # Init model (ensure BiLSTM)
+    # Model setup
     model = OCRModel(num_classes=len(CHAR2IDX)).to(device)
     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)
 
     for epoch in range(epochs):
-        # Load new dataset for current curriculum stage
         dataset = get_dataset_for_epoch(epoch)
         dataloader = DataLoader(dataset, batch_size=16, shuffle=True, collate_fn=custom_collate_fn)
 
         model.train()
         running_loss = 0.0
 
-        # CTC warmup: reduced LR during initial epochs
+        # Warmup learning rate
         if epoch < 5:
             warmup_lr = learning_rate * 0.2
             for param_group in optimizer.param_groups:
@@ -230,12 +240,12 @@ def train_model(font_file, epochs=100, learning_rate=0.001):
             targets = targets.to(device)
             target_lengths = target_lengths.to(device)
 
-            output = model(img)  # [B, T, C]
+            output = model(img)
             seq_len = output.size(1)
             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]
+            log_probs = output.log_softmax(2).transpose(0, 1)
             loss = criterion(log_probs, targets, input_lengths, target_lengths)
 
             optimizer.zero_grad()
@@ -248,13 +258,15 @@ def train_model(font_file, epochs=100, learning_rate=0.001):
         scheduler.step(avg_loss)
         print(f"[{epoch + 1}/{epochs}] Loss: {avg_loss:.4f}")
 
-    # Save the model
+    # Save the model to ./models
     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)
+    model_path = os.path.join("./models", model_name)
+    save_model(model, model_path)
+
     ocr_model = model
+    return f"✅ Training complete! Model saved as '{model_path}'"
 
-    return f"✅ Training complete! Model saved as '{model_name}'"
 
 
 
@@ -376,11 +388,11 @@ def generate_labels(font_file=None, num_labels: int = 25):
     global font_path
 
     try:
-        if font_file:
-            font_path = "./temp_font_labels.ttf"
-            with open(font_file.name, "rb") as uploaded:
-                with open(font_path, "wb") as f:
-                    f.write(uploaded.read())
+        if font_file and font_file != "None":
+            font_path = os.path.abspath(font_file)
+        else:
+            font_path = None
+
         if font_path is None or not os.path.exists(font_path):
             font = ImageFont.load_default()
         else:
@@ -391,7 +403,6 @@ def generate_labels(font_file=None, num_labels: int = 25):
         images = []
 
         for label in labels:
-            # Measure text size and calculate padded image dimensions
             bbox = font.getbbox(label)
             text_w = bbox[2] - bbox[0]
             text_h = bbox[3] - bbox[1]
@@ -399,12 +410,10 @@ def generate_labels(font_file=None, num_labels: int = 25):
             img_w = text_w + pad * 2
             img_h = text_h + pad * 2
 
-            # Create image and draw text
             img = Image.new("L", (img_w, img_h), color=255)
             draw = ImageDraw.Draw(img)
             draw.text((pad, pad), label, font=font, fill=0)
 
-            # Save to ./labels/sanitized_label/timestamp.png
             safe_label = sanitize_filename(label)
             timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
             label_dir = os.path.join("./labels", safe_label)
@@ -424,6 +433,10 @@ def generate_labels(font_file=None, num_labels: int = 25):
         draw.text((10, 50), f"Error: {str(e)}", fill=(255, 0, 0))
         return [error_img]
 
+def list_fonts():
+    fonts = [f for f in os.listdir() if f.lower().endswith((".ttf", ".otf"))]
+    return ["None"] + fonts if fonts else ["None"]
+
 custom_css = """
 #label-gallery .gallery-item img {
     height: 43px;           /* 32pt ≈ 43px */
@@ -444,7 +457,7 @@ custom_css = """
 
 # --------- Updated Gradio UI with new tab --------- #
 with gr.Blocks(css=custom_css) as demo:
-    with gr.Tab("1. Upload Font & Train"):
+    with gr.Tab("【Train OCR Model】"):
         font_file = gr.File(label="Upload .ttf or .otf font", file_types=[".ttf", ".otf"])
         epochs_input = gr.Slider(minimum=1, maximum=4096, value=256, step=1, label="Epochs")
         lr_input = gr.Slider(minimum=0.001, maximum=0.1, value=0.05, step=0.001, label="Learning Rate")
@@ -453,8 +466,28 @@ with gr.Blocks(css=custom_css) as demo:
 
     train_button.click(fn=train_model, inputs=[font_file, epochs_input, lr_input], outputs=train_status)
 
+    with gr.Tab("【Generate Labels】"):
+        font_file_labels = gr.Dropdown(
+            choices=list_fonts(),
+            label="Optional font for label image",
+            interactive=True,
+        )
+        num_labels = gr.Number(value=20, label="Number of labels to generate", precision=0, interactive=True)
+        gen_button = gr.Button("Generate Label Grid")
+
+        gen_button.click(
+            fn=generate_labels,
+            inputs=[font_file_labels, num_labels],
+            outputs=gr.Gallery(
+                label="Generated Labels",
+                columns=16,                  # 16 tiles per row
+                object_fit="contain",       # Maintain aspect ratio
+                height="100%",              # Allow full app height
+                elem_id="label-gallery"     # For CSS targeting
+            )
 
-    with gr.Tab("2. Use Trained Model"):
+        )
+    with gr.Tab("【Recognize Text】"):
         model_list = gr.Dropdown(choices=list_saved_models(), label="Select OCR Model")
         refresh_btn = gr.Button("🔄 Refresh Models")
         load_model_btn = gr.Button("Load Model")  # <-- new button
@@ -472,23 +505,7 @@ with gr.Blocks(css=custom_css) as demo:
 
         predict_btn.click(fn=predict_text, inputs=image_input, outputs=output_text)
 
-    with gr.Tab("3. Generate Labels"):
-        font_file_labels = gr.File(label="Optional font for label image", file_types=[".ttf", ".otf"])
-        num_labels = gr.Number(value=20, label="Number of labels to generate", precision=0, interactive=True)
-        gen_button = gr.Button("Generate Label Grid")
-
-        gen_button.click(
-            fn=generate_labels,
-            inputs=[font_file_labels, num_labels],
-            outputs=gr.Gallery(
-                label="Generated Labels",
-                columns=16,                  # 16 tiles per row
-                object_fit="contain",       # Maintain aspect ratio
-                height="100%",              # Allow full app height
-                elem_id="label-gallery"     # For CSS targeting
-            )
 
-        )