commit 6b730808c5a7928af4d4bd2e79378a1a63cc0aed Author: Romans Juškevičs Date: Fri Sep 18 08:51:00 2026 +0000 Upload files to "/" diff --git a/README.md b/README.md new file mode 100644 index 0000000..cd6d28e --- /dev/null +++ b/README.md @@ -0,0 +1,80 @@ +# Latviešu rokrakstu atpazīšana — sākotnējā versija + +Sākotnējā (baseline) implementācija pēc tavas blokshēmas. **AI bloks blokshēmā bija uzzīmēts nepareizi** — CNN un BiLSTM tur izskatās pēc diviem neatkarīgiem zariem, bet praksē tā strādāt nevar. Reālajā arhitektūrā (CRNN, standarts rokraksta/OCR atpazīšanai) tie ir **secīgi** posmi + trūkst CTC slāņa, kas savieno tīkla izvadi ar mainīga garuma tekstu: + +``` +attēls -> CNN (vizuālās pazīmes, "Līnijas atpazīšana" / "Pareiza novietošana") + -> BiLSTM x2 (konteksts pa sekvenci, "Vārdu atpazīšana" / "Pieturzīmes" / "Cipari") + -> Linear + CTC (blokshēmā vispār nebija — bez tā nevar mainīga garuma + izvadi savienot ar tekstu bez rakstzīmju līmeņa segmentācijas) +``` + +CTC arī nozīmē, ka "Teksta zonas atpazīšana????" (ar jautājuma zīmēm blokshēmā — +acīmredzot arī tev bija šaubas) lielākoties nav vajadzīga uz ievades puses: +var padot veselas rindas attēlu, nevis atsevišķus burtus. + +## Faili + +| Fails | Blokshēmas daļa | +|---|---| +| `alphabet.py` | rakstzīmju kopa (ā č ē ģ ī ķ ļ ņ š ū ž u.c.) + CTC kodēšana/dekodēšana | +| `synthetic_data.py` | "Sintētiskie dati" — ģenerē attēlus no teksta, ar rotāciju/šķības/troksni | +| `dataset.py` | "Reāli dati" (`RealHTRDataset`) + sintētiskie (`SyntheticHTRDataset`) | +| `model.py` | laboti "AI" — CRNN (CNN -> BiLSTM -> CTC) | +| `train.py` | apmācības cikls | +| `infer.py` | "OUTPUT" — attēls -> teksts | + +## Palaišana + +```bash +pip install -r requirements.txt + +# ātrs tests — ģenerē pāris paraugattēlus +python synthetic_data.py + +# apmācība tikai uz sintētiskajiem datiem (pagaidu risinājums, kamēr nav reālu skenējumu) +python train.py --steps 2000 --batch-size 32 + +# kad būs reāli dati (skat. zemāk), pievieno tos apmācībai +python train.py --steps 5000 --real-data data/real_dataset + +# secinājums uz viena attēla +python infer.py --checkpoint checkpoints/crnn_step2000.pt --image samples/sample_0.png +``` + +## Kā pievienot reālus datus + +Sagatavo mapi (tas ir blokshēmas "Rokrakstu bāzes sagatavošana" solis): + +``` +data/real_dataset/ + labels.csv # kolonnas: filename,text + images/ + 0001.png + 0002.png +``` + +`labels.csv` piemērs: +```csv +filename,text +0001.png,Labdien! Kā jums iet? +0002.png,Rīgā šodien līst. +``` + +## Kas šobrīd ir vienkāršots / jāuzlabo tālāk + +- **Sintētiskie dati** izmanto parastu (ne-rokraksta) fontu (DejaVu Sans) ar + nelielām deformācijām. Reālam rokrakstam ieteicams iemest `fonts/` mapē + kursīvu/rokraksta stila `.ttf` fontu ar latviešu diakritiku atbalstu — + ģenerators to automātiski izmantos. +- Nav vēl datu augmentācijas variāciju pēc reāliem paraugiem (piem., papīra + tekstūras, tintes izplūduma) — der pievienot, tiklīdz ir reāli skeni, lai + redzētu, kāda veida troksnis tos raksturo. +- "Cik piemēru?" no blokshēmas — nav fiksēts skaitlis kodā; `--steps` un + `--synth-per-epoch` kontrolē apjomu, pielāgo pēc vajadzības. +- Šī vide nespēja lokāli instalēt/palaist pilnu PyTorch (CUDA atkarības prasa + vairāk diska vietas nekā šeit pieejams), tāpēc modeļa forward/apmācības + kods nav palaists end-to-end šajā sandbox — sintaktiski pārbaudīts un + arhitektūra ir standarta, pārbaudīta CRNN+CTC shēma, bet ieteicams pirmo + palaišanu izdarīt ar mazu `--steps` skaitli, lai pārliecinātos, ka viss + strādā tavā vidē. diff --git a/alphabet.py b/alphabet.py new file mode 100644 index 0000000..5f27291 --- /dev/null +++ b/alphabet.py @@ -0,0 +1,40 @@ +""" +Charset / vocabulary for Latvian handwriting recognition (CTC-based). + +Index 0 is reserved for the CTC "blank" symbol - do not remove it. +""" + +LATVIAN_LETTERS = "aābcčdeēfgģhiījkķlļmnņoprsštuūvzž" +BASE_LATIN = "qwxy" # appear in loanwords / foreign names, kept for robustness +DIGITS = "0123456789" +PUNCT = " .,!?-:;()\"'/" + +CHARS = sorted(set(LATVIAN_LETTERS + LATVIAN_LETTERS.upper() + BASE_LATIN + BASE_LATIN.upper() + DIGITS + PUNCT)) + +BLANK = "" +ALPHABET = [BLANK] + CHARS + +CHAR_TO_IDX = {c: i for i, c in enumerate(ALPHABET)} +IDX_TO_CHAR = {i: c for i, c in enumerate(ALPHABET)} + +NUM_CLASSES = len(ALPHABET) + + +def encode(text: str) -> list[int]: + """Text -> list of class indices (no blanks inserted, CTCLoss wants targets without blanks).""" + unknown = set(ch for ch in text if ch not in CHAR_TO_IDX) + if unknown: + raise ValueError(f"Unsupported characters in text: {unknown!r}. Extend alphabet.py CHARS.") + return [CHAR_TO_IDX[ch] for ch in text] + + +def decode_greedy(indices: list[int]) -> str: + """Collapse repeats and drop blanks - standard CTC greedy decoding.""" + out = [] + prev = None + for idx in indices: + if idx != prev: + if idx != 0: # 0 == blank + out.append(IDX_TO_CHAR[idx]) + prev = idx + return "".join(out) diff --git a/dataset.py b/dataset.py new file mode 100644 index 0000000..1996fff --- /dev/null +++ b/dataset.py @@ -0,0 +1,100 @@ +""" +Dataset classes. + +- SyntheticHTRDataset: infinite on-the-fly synthetic samples (current stand-in, + "Sintētiskie dati" on the flowchart). +- RealHTRDataset: loads scanned line images once you have them + ("Reāli dati" on the flowchart). Expects a folder with images plus a + labels.csv of `filename,text` pairs - this is the "Rokrakstu bāzes + sagatavošana" step from the diagram, done offline before training. + +Both return (image_tensor[1,H,W], text_string). Collation pads variable-width +images to the batch max width, since line images differ in length. +""" + +import csv +import os + +import torch +from PIL import Image +from torch.utils.data import Dataset + +from synthetic_data import TARGET_HEIGHT, generate_batch, load_corpus, render_text_line + +import numpy as np + + +def image_to_tensor(img: Image.Image) -> torch.Tensor: + arr = np.array(img.convert("L"), dtype=np.float32) / 255.0 + arr = 1.0 - arr # invert: background 0, ink ~1 (easier for the CNN) + return torch.from_numpy(arr).unsqueeze(0) # [1, H, W] + + +class SyntheticHTRDataset(Dataset): + """Generates `length` synthetic (image, text) pairs per epoch, freshly each time.""" + + def __init__(self, length: int = 2000, corpus: list[str] | None = None): + self.length = length + self.corpus = corpus or load_corpus() + + def __len__(self): + return self.length + + def __getitem__(self, idx): + import random + + text = random.choice(self.corpus) + img = render_text_line(text) + return image_to_tensor(img), text + + +class RealHTRDataset(Dataset): + """ + Loads real scanned handwriting samples. + + Expected layout: + root/ + labels.csv # header: filename,text + images/ + 0001.png + 0002.png + ... + """ + + def __init__(self, root: str): + self.root = root + self.samples: list[tuple[str, str]] = [] + labels_path = os.path.join(root, "labels.csv") + if not os.path.exists(labels_path): + raise FileNotFoundError( + f"Expected {labels_path} with columns 'filename,text'. " + "This is the 'Rokrakstu bāzes sagatavošana' step - prepare it offline." + ) + with open(labels_path, encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + self.samples.append((row["filename"], row["text"])) + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + filename, text = self.samples[idx] + img = Image.open(os.path.join(self.root, "images", filename)).convert("L") + w, h = img.size + new_w = max(1, int(w * (TARGET_HEIGHT / h))) + img = img.resize((new_w, TARGET_HEIGHT), Image.BILINEAR) + return image_to_tensor(img), text + + +def collate_batch(batch): + """Pad images to max width in batch; return images, texts, and original widths.""" + imgs, texts = zip(*batch) + max_w = max(img.shape[-1] for img in imgs) + padded = torch.zeros(len(imgs), 1, TARGET_HEIGHT, max_w) + widths = torch.zeros(len(imgs), dtype=torch.long) + for i, img in enumerate(imgs): + w = img.shape[-1] + padded[i, :, :, :w] = img + widths[i] = w + return padded, list(texts), widths diff --git a/infer.py b/infer.py new file mode 100644 index 0000000..44e41d9 --- /dev/null +++ b/infer.py @@ -0,0 +1,51 @@ +""" +Inference. + +Usage: + python infer.py --checkpoint checkpoints/crnn_step2000.pt --image samples/sample_0.png + +Maps to the flowchart's OUTPUT branch: "Teksts digitālajā veidā" -> +"Koeficientu pārveidošana burtos" (= CTC decode, logits -> characters) -> +"Teksta rakstīšana" (= printing/returning the final string). +""" + +import argparse + +import torch +from PIL import Image + +from alphabet import NUM_CLASSES, decode_greedy +from dataset import image_to_tensor +from model import CRNN +from synthetic_data import TARGET_HEIGHT + + +def load_image(path: str) -> torch.Tensor: + img = Image.open(path).convert("L") + w, h = img.size + new_w = max(1, int(w * (TARGET_HEIGHT / h))) + img = img.resize((new_w, TARGET_HEIGHT), Image.BILINEAR) + return image_to_tensor(img).unsqueeze(0) # [1,1,H,W] + + +def predict(model: CRNN, image_tensor: torch.Tensor, device: torch.device) -> str: + model.eval() + with torch.no_grad(): + log_probs = model(image_tensor.to(device)) # [T,1,C] + pred_indices = log_probs.argmax(dim=2).squeeze(1).tolist() # [T] + return decode_greedy(pred_indices) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=str, required=True) + parser.add_argument("--image", type=str, required=True) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = CRNN(num_classes=NUM_CLASSES).to(device) + model.load_state_dict(torch.load(args.checkpoint, map_location=device)) + + image_tensor = load_image(args.image) + text = predict(model, image_tensor, device) + print(f"Recognized text: {text!r}") diff --git a/model.py b/model.py new file mode 100644 index 0000000..ec028f1 --- /dev/null +++ b/model.py @@ -0,0 +1,82 @@ +""" +Model architecture - this is the corrected version of the flowchart's "AI" box. + +On the diagram, CNN and BiLSTM were drawn as two independent branches off "AI". +For text recognition that's not how it works: they're sequential stages of one +pipeline (a CRNN), plus a CTC layer the diagram was missing entirely: + + image -> CNN (visual features, replaces "Līnijas atpazišana" / + "Pareiza novietošana" - i.e. stroke patterns + spatial layout) + -> reshape rows->sequence + -> BiLSTM x2 (context over the sequence, replaces "Vārdu atpazīšana" / + "Pieturzīmes" / "Cipari" - i.e. word/punctuation/digit recognition + needs context from neighboring characters, which only the LSTM + stage gives you) + -> Linear classifier over the alphabet + -> CTC loss/decoding (aligns variable-length predictions to text + without needing per-character bounding boxes) + +CTC is what makes "Teksta zonas atpazīšana????" mostly unnecessary for the +input side too: you don't need to segment characters up front, just feed +whole line images in. +""" + +import torch +import torch.nn as nn + + +class CNNBackbone(nn.Module): + """Reduces a [B,1,32,W] line image to a [B,512,1,W'] feature map.""" + + def __init__(self): + super().__init__() + self.net = nn.Sequential( + nn.Conv2d(1, 64, 3, 1, 1), nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), # 32x W -> 16 x W/2 + nn.Conv2d(64, 128, 3, 1, 1), nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), # -> 8 x W/4 + nn.Conv2d(128, 256, 3, 1, 1), nn.ReLU(inplace=True), + nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(inplace=True), nn.MaxPool2d((2, 1), (2, 1)), # -> 4 x W/4 + nn.Conv2d(256, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), + nn.Conv2d(512, 512, 3, 1, 1), nn.BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d((2, 1), (2, 1)), # -> 2 x W/4 + nn.Conv2d(512, 512, 2, 1, 0), nn.ReLU(inplace=True), # -> 1 x (W/4 - 1) + ) + + def forward(self, x): + return self.net(x) + + +class BiLSTMHead(nn.Module): + def __init__(self, in_dim: int, hidden: int, num_classes: int, num_layers: int = 2): + super().__init__() + self.lstm = nn.LSTM( + in_dim, hidden, num_layers=num_layers, bidirectional=True, batch_first=False + ) + self.fc = nn.Linear(hidden * 2, num_classes) + + def forward(self, x): + # x: [T, B, in_dim] + out, _ = self.lstm(x) + return self.fc(out) # [T, B, num_classes] + + +class CRNN(nn.Module): + def __init__(self, num_classes: int, lstm_hidden: int = 256): + super().__init__() + self.cnn = CNNBackbone() + self.rnn = BiLSTMHead(in_dim=512, hidden=lstm_hidden, num_classes=num_classes) + + def forward(self, images: torch.Tensor) -> torch.Tensor: + """ + images: [B, 1, 32, W] + returns log-probs: [T, B, num_classes], T = sequence length after CNN downsampling + """ + feats = self.cnn(images) # [B, 512, 1, W'] + feats = feats.squeeze(2) # [B, 512, W'] + feats = feats.permute(2, 0, 1) # [W'(=T), B, 512] + logits = self.rnn(feats) # [T, B, num_classes] + return logits.log_softmax(dim=2) + + def output_length(self, input_width: int) -> int: + """Sequence length T the CNN produces for a given input image width (for CTC input_lengths).""" + w = input_width // 2 // 2 # two stride-2 pools + w = w - 1 # final kernel=2,stride=1,pad=0 conv + return max(w, 1)