Update app.py

app.py

@@ -41,28 +41,42 @@ decoder = BeamSearchDecoderCTC(alphabet)
 
 # --------- Dataset --------- #
 class OCRDataset(Dataset):
-    def __init__(self, font_path, size=1000):
+    def __init__(self, font_path, size=1000, label_length_range=(4, 7)):
         self.font = ImageFont.truetype(font_path, 32)
-        self.samples = ["".join(np.random.choice(list(CHARS), np.random.randint(4, 7)))
-                        for _ in range(size)]
+        self.label_length_range = label_length_range
+        self.samples = [
+            "".join(np.random.choice(list(CHARS), np.random.randint(*self.label_length_range)))
+            for _ in range(size)
+        ]
 
         self.transform = transforms.Compose([
-            transforms.Grayscale(),
+            transforms.ToTensor(),  # must be first
+            transforms.Normalize((0.5,), (0.5,)),
             transforms.Resize((IMAGE_HEIGHT, IMAGE_WIDTH)),
-            transforms.ToTensor(),
-            transforms.Normalize((0.5,), (0.5,))
+            transforms.RandomApply([transforms.GaussianBlur(kernel_size=3)], p=0.3),
+            transforms.RandomApply([transforms.RandomAffine(degrees=10, translate=(0.1, 0.1))], p=0.3),
         ])
-
     def __len__(self):
         return len(self.samples)
 
     def __getitem__(self, idx):
-        text = self.samples[idx]
-        img = self.render_text(text)
-        img = self.transform(img)  # convert PIL to tensor with normalization
+        label = self.samples[idx]
+
+        # Create an image with padding
+        pad = 8
+        w = self.font.getlength(label)
+        h = self.font.size
+        img_w, img_h = int(w + 2 * pad), int(h + 2 * pad)
+        img = Image.new("L", (img_w, img_h), 255)
+        draw = ImageDraw.Draw(img)
+        draw.text((pad, pad), label, font=self.font, fill=0)
+
+        img = self.transform(img)
+        label_encoded = torch.tensor([CHAR2IDX[c] for c in label], dtype=torch.long)
+        label_length = torch.tensor(len(label_encoded), dtype=torch.long)
+
+        return img, label_encoded, label_length
 
-        label = torch.tensor([CHAR2IDX[c] for c in text], dtype=torch.long)
-        return img, label
 
 
     def render_text(self, text):
@@ -125,7 +139,7 @@ def greedy_decode(log_probs):
 
 # --------- Custom Collate --------- #
 def custom_collate_fn(batch):
-    images, labels = zip(*batch)
+    images, labels, _ = zip(*batch)
     images = torch.stack(images, 0)
 
     flat_labels = []
@@ -163,34 +177,54 @@ def train_model(font_file, epochs=100, learning_rate=0.001):
     import time
     global font_path, ocr_model
 
-    # Save the uploaded font file
+    # Save uploaded font
     font_name = os.path.splitext(os.path.basename(font_file.name))[0]
     font_path = f"./{font_name}.ttf"
     with open(font_file.name, "rb") as uploaded:
         with open(font_path, "wb") as f:
             f.write(uploaded.read())
 
-    # Load dataset
-    dataset = OCRDataset(font_path)
-    dataloader = DataLoader(dataset, batch_size=16, shuffle=True, collate_fn=custom_collate_fn)
+    # Curriculum learning: Start with shorter labels, increase over time
+    def get_dataset_for_epoch(epoch):
+        if epoch < epochs // 3:
+            label_len = (3, 4)
+        elif epoch < 2 * epochs // 3:
+            label_len = (4, 6)
+        else:
+            label_len = (5, 7)
+        return OCRDataset(font_path, label_length_range=label_len)
+
+    # Visualize one sample from initial dataset
+    dataset = get_dataset_for_epoch(0)
+    img, label, _ = dataset[0]  # Ignore the 3rd value (e.g., label length)
 
-    # 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, loss, optimizer, scheduler
+    # Init model (ensure BiLSTM)
     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)
 
-    # Training loop
     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
+        if epoch < 5:
+            warmup_lr = learning_rate * 0.2
+            for param_group in optimizer.param_groups:
+                param_group['lr'] = warmup_lr
+        else:
+            for param_group in optimizer.param_groups:
+                param_group['lr'] = learning_rate
+
         for img, targets, target_lengths in dataloader:
             img = img.to(device)
             targets = targets.to(device)
@@ -214,7 +248,7 @@ 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 trained model
+    # Save the 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)
@@ -226,6 +260,7 @@ def train_model(font_file, epochs=100, learning_rate=0.001):
 
 
 
+
 def preprocess_image(image: Image.Image):
     img_cv = np.array(image.convert("L"))
 
@@ -300,7 +335,11 @@ def predict_text(image: Image.Image, ground_truth: str = None, debug: bool = Fal
         output = ocr_model(img_tensor)           # (1, T, C)
         log_probs = output.log_softmax(2)[0]     # (T, C)
 
-        pred_text = decoder.decode(log_probs.cpu().numpy())  # Best beam path
+        # Decode best beam path (string)
+        pred_text_raw = decoder.decode(log_probs.cpu().numpy())
+        pred_chars = pred_text_raw.replace("<BLANK>", "")
+        # Remove <BLANK> tokens if present (assuming <BLANK> is in vocab)
+        pred_text = ''.join([c for c in pred_chars if c != "<BLANK>"])
 
         # Confidence: mean max prob per timestep
         probs = log_probs.exp()
@@ -322,7 +361,8 @@ def predict_text(image: Image.Image, ground_truth: str = None, debug: bool = Fal
         if ground_truth:
             print("Ground Truth:", ground_truth)
 
-    return f"<strong>Prediction:</strong> {pretty_output}<br><strong>Confidence:</strong> {avg_conf:.2%}{sim_score}"
+    return f"<strong>Prediction:</strong> <strong>{pretty_output}</strong><br><strong>Confidence:</strong> {avg_conf:.2%}{sim_score}"
+
 
 
 # New helper function: generate label images grid