source code

src/Text_Localization/prepare_dataset.py

@@ -0,0 +1,177 @@
+import numpy as np
+import os
+import matplotlib.pyplot as plt
+import cv2
+import xml.etree.ElementTree as ET
+import shutil
+import yaml
+
+from sklearn.model_selection import train_test_split
+
+location_path = r'Dataset/locations.xml'
+tree = ET.parse(location_path)
+root = tree.getroot()
+
+
+def get_gt_bboxes(location_path):
+    """get all the gt bbox of text in dataset
+
+    Args:
+        location_path: (path)
+    Return:
+        gt_imagepaths[1] (list): image's name
+        gt_locations (list): bboxes of each image
+    """
+    gt_imagepaths = []
+    gt_imagesizes = []
+    gt_locations = []
+    
+    for image in root:
+        # get path to image
+        image_name = image[0].text
+        image_path = os.path.join('Dataset', image_name)
+        gt_imagepaths.append(image_path)
+        
+        # get the image size
+        w = image[1].get('x')
+        h = image[1].get('y')
+        gt_imagesizes.append([h, w])
+        
+        # bboxes in the image
+        bbs = []
+        for bbox in image[2]:
+            x = np.int64(float(bbox.get('x')))
+            y = np.int64(float(bbox.get('y')))
+            width = np.int64(float(bbox.get('width')))
+            height = np.int64(float(bbox.get('height')))
+            bbs.append([x, y, width, height])
+            
+        gt_locations.append(bbs)
+    
+    return gt_imagepaths, gt_imagesizes, gt_locations
+    
+gt_imagepaths, gt_imagesizes, gt_locations = get_gt_bboxes(location_path)
+
+def visualize_gt_bboxes(image_path, gt_locations):
+    image = cv2.imread(image_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    for gt_location in gt_locations:
+        x, y, width, height = gt_location
+        x, y, width, height = int(x), int(y), int(width), int(height)
+        
+        image = cv2.rectangle(image, (x, y), (x+width, y+height), color=(255, 0, 0), thickness=2)
+        
+    plt.imshow(image)
+    plt.axis('off')
+    plt.show()
+    
+def visualize_gt_bboxes_yolo(image_path, gt_location_yolo):
+    image = cv2.imread(image_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    image_height, image_width = image.shape[:2]
+    
+    # Convert to original format
+    for data in gt_location_yolo:
+        xc, yc, w, h = data[1:]
+        xmin = int((xc - w/2) * image_width)
+        ymin = int((yc - h/2) * image_height)
+        xmax = int((xc + w/2) * image_width)
+        ymax = int((yc + h/2) * image_height)
+        
+        image = cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color=(255, 0, 0), thickness=2)
+        
+    plt.imshow(image)
+    plt.axis('off')
+    plt.show()
+    
+    
+def convert_yolo_format(gt_locations, gt_imagesizes):
+    gt_locations_yolo = []
+    
+    for image, image_size in zip(gt_locations, gt_imagesizes):
+        gt_location_yolo = []
+        for gt_location in image:
+            x, y, w, h = gt_location
+            image_height, image_width = image_size
+            
+            xc = (x + w/2) / float(image_width)
+            yc = (y + h/2) / float(image_height)
+            width = w / float(image_width)
+            height = h / float(image_height)
+            
+            # class = 0 -> meaning contains text
+            class_id = 0
+            gt_location_yolo.append([class_id, xc, yc, width, height])
+        
+        gt_locations_yolo.append(gt_location_yolo)
+        
+    return gt_locations_yolo
+        
+gt_locations_yolo = convert_yolo_format(gt_locations, gt_imagesizes)
+
+def save_data_into_yolo_folder(data, src_img_dir, save_dir):
+    # Create folder if not exist
+    os.makedirs(save_dir, exist_ok=True)
+    
+    # Make images and labels folder
+    os.makedirs(os.path.join(save_dir, 'images'), exist_ok=True)
+    os.makedirs(os.path.join(save_dir, 'labels'), exist_ok=True)
+    
+    # write data into yolo folder
+    for dt in data:
+        # copy data
+        image_path = dt[0]
+        shutil.copy(image_path, os.path.join(save_dir, 'images'))
+        
+        #copy labels
+        image_name = os.path.basename(image_path)
+        image_name = os.path.splitext(image_name)[0]
+        
+        with open(os.path.join(save_dir, 'labels', f'{image_name}.txt'), "w") as f:
+            for label in dt[1]:
+                label_str = " ".join(map(str, label))
+                f.write(f'{label_str}\n')
+
+
+
+seed = 0
+val_size = 0.2
+test_size = 0.125
+dataset = [[gt_imagepath, gt_location_yolo] for gt_imagepath, gt_location_yolo in zip(gt_imagepaths, gt_locations_yolo)]
+train_data, val_data = train_test_split(dataset, test_size=val_size, random_state=42, shuffle=True)
+train_data, test_data = train_test_split(train_data, test_size=test_size, random_state=42, shuffle=True)
+
+save_yolo_data_dir = 'yolo_data'
+os.makedirs(save_yolo_data_dir, exist_ok=True)
+save_data_into_yolo_folder(
+    data=train_data,
+    src_img_dir=save_yolo_data_dir,
+    save_dir=os.path.join(save_yolo_data_dir, 'train')
+)
+save_data_into_yolo_folder(
+    data=val_data,
+    src_img_dir=save_yolo_data_dir,
+    save_dir=os.path.join(save_yolo_data_dir, 'val')
+)
+save_data_into_yolo_folder(
+    data=test_data,
+    src_img_dir=save_yolo_data_dir,
+    save_dir=os.path.join(save_yolo_data_dir, 'test')
+)
+
+class_label = ['text']
+# Create data.yaml file
+data_yaml = {
+    "path": '../yolo_data',
+    'train': 'train/images',
+    'test': 'test/images',
+    'val': 'val/images',
+    'nc': 1,
+    'names': class_label
+}
+
+yolo_yaml_path = os.path.join(save_yolo_data_dir, 'data.yaml')
+with open(yolo_yaml_path, "w") as f:
+    yaml.dump(data_yaml, f, default_flow_style=False)

src/Text_Localization/text_localization.py

@@ -0,0 +1,33 @@
+import ultralytics
+import json
+from ultralytics import YOLO
+
+yolo_yaml_path = 'yolo_data/data.yaml'
+config_path = 'src/config.json'
+
+def load_json_config(config_path):
+    with open(config_path, "r") as f:
+        config = json.load(f)
+        
+    return config
+
+config = load_json_config(config_path)
+
+# Load model
+model = YOLO('yolo11m.pt')
+
+# Train model
+results = model.train(
+    data=yolo_yaml_path,
+    epochs=config['yolov11']['epochs'],
+    imgsz=config['yolov11']['image_size'],
+    cache=config['yolov11']['cache'],
+    patience=config['yolov11']['patience'],
+    plots=config['yolov11']['plots']
+)
+
+# Evaluate model
+model_path = 'checkpoints/yolov11m.pt'
+model = YOLO(model_path)
+
+metrics = model.val()

src/Text_Recognization/dataloader.py

@@ -0,0 +1,96 @@
+import torch
+import torchvision
+import json
+import sys
+
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from sklearn.model_selection import train_test_split
+from src.Text_Recognization.prepare_dataset import *
+
+# data augmentation
+data_transforms = {
+    "train": transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Resize((100, 400)),
+            transforms.ColorJitter(
+                brightness=0.5,
+                contrast=0.5,
+                saturation=0.5
+            ),
+            transforms.GaussianBlur(3),
+            transforms.RandomAffine(
+                degrees=1,
+                shear=1
+            ),
+            transforms.RandomPerspective(
+                distortion_scale=0.3,
+                p=0.5
+            ),
+            transforms.RandomRotation(degrees=15),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ]
+    ),
+    "val": transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Resize((100, 400)),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ]
+    )
+}
+
+def load_json_config(config_path):
+    with open(config_path, "r") as f:
+        config = json.load(f)
+        
+    return config
+
+# Dataloader
+class STRDataset(Dataset):
+    def __init__(self, image_paths, labels, char_to_idx, transforms=None):
+        self.image_paths = image_paths
+        self.labels = labels
+        self.char_to_idx = char_to_idx
+        self.transforms= transforms
+        
+    def __len__(self):
+        return len(self.image_paths)
+    
+    def __getitem__(self, idx):
+        image = cv2.imread(self.image_paths[idx])
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        
+        if self.transforms:
+            image = self.transforms(image)
+        
+        label_encoded, length = encode(self.labels[idx], self.char_to_idx, self.labels)
+        
+        return image, label_encoded, length
+    
+def get_dataloader():
+    val_size = 0.1
+    test_size = 0.1
+    root_path = 'Dataset'
+    config_path = 'src/config.json'
+
+    # get image paths and labels
+    image_paths, labels = get_imagepaths_and_labels(root_path)
+    char_to_idx, idx_to_char = build_vocab(root_path)
+
+
+    config = load_json_config(config_path)
+
+    X_train, X_val, y_train, y_val = train_test_split(image_paths, labels, test_size=val_size, random_state=42, shuffle=True)
+    X_train, X_test, y_train, y_test = train_test_split(X_train, y_train, test_size=test_size, random_state=42, shuffle=True)
+    train_dataset = STRDataset(X_train, y_train, char_to_idx, transforms=data_transforms['train'])
+    train_loader = DataLoader(train_dataset, batch_size=config['CRNN']['batch_size'], shuffle=True)
+
+    val_dataset = STRDataset(X_val, y_val, char_to_idx, transforms=data_transforms['val'])
+    val_loader = DataLoader(val_dataset, batch_size=config['CRNN']['batch_size'], shuffle=True)
+
+    test_dataset = STRDataset(X_test, y_test, char_to_idx, transforms=data_transforms['val'])
+    test_loader = DataLoader(test_dataset, batch_size=config['CRNN']['batch_size'], shuffle=True)
+    
+    return train_loader, val_loader, test_loader

src/Text_Recognization/prepare_dataset.py

@@ -0,0 +1,221 @@
+import os
+import random
+import numpy as np
+import xml.etree.ElementTree as ET
+import cv2
+import json
+import matplotlib.pyplot as plt
+import argparse
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def extract_data_from_xml(root_path):
+    words_path = os.path.join(root_path, 'words.xml')
+    tree = ET.parse(words_path)
+    root = tree.getroot()
+    
+    image_paths = []
+    image_sizes = []
+    image_labels = []
+    bboxes = []
+    
+    for image in root:
+        imagename = image[0].text
+        image_path = os.path.join(root_path, imagename)
+        image_paths.append(image_path)
+        
+        image_height = image[1].get('x')
+        image_width = image[1].get('y')
+        image_sizes.append([image_height, image_width])
+        
+        bboxes_in_image = []
+        labels_in_bboxes = []
+        for bbox in image[2]:
+            x = float(bbox.get('x'))
+            y = float(bbox.get('y'))
+            width = float(bbox.get('width'))
+            height = float(bbox.get('height'))
+            bboxes_in_image.append([x, y, width, height])
+            
+            # get text in this bbox
+            labels = bbox.find('tag').text
+            labels_in_bboxes.append(labels)
+            
+        bboxes.append(bboxes_in_image)
+        image_labels.append(labels_in_bboxes)
+        
+    return image_paths, image_sizes, bboxes, image_labels
+
+def visualize_gt_bboxes(image_path, gt_locations, gt_labels):
+    image = cv2.imread(image_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    for gt_location, gt_label in zip(gt_locations, gt_labels):
+        x, y, width, height = gt_location
+        x, y, width, height = int(x), int(y), int(width), int(height)
+        
+        image = cv2.rectangle(image, (x, y), (x+width, y+height), color=(255, 0, 0), thickness=2)
+        image = cv2.putText(image, gt_label, (x, y-10), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale = 3, color=(255, 0, 0), thickness=2)
+        
+    plt.imshow(image)
+    plt.axis('off')
+    plt.show()
+
+
+def split_bboxes_from_image(image_paths, image_labels, bboxes, save_dir):
+    """create a new dataset contains bboxes and corresponding labels
+
+    Args:
+        image_paths
+        image_labels
+        bboxes
+        save_dir
+    Return:
+        non-return
+    """
+    os.makedirs(save_dir, exist_ok=True)
+    os.makedirs('unvalid_images', exist_ok=True)
+    
+    bboxes_idx = 0
+    unvalid_bboxes = 0
+    new_labels = []         # List to store labels
+    for image_path, bbox, label in zip(image_paths, bboxes, image_labels):
+        image = cv2.imread(image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        
+        if image is None:
+            print(image_path)
+            continue
+        
+        for bb, lb in zip(bbox, label):
+            x, y, width, height = bb
+            x, y, width, height = int(x), int(y), int(width), int(height)
+            
+            cropped_text = image[y:y+height, x:x+width]
+            
+            # Filter if x, y, width, height is invalid cordinates
+            if x < 0 or y < 0 or width < 0 or height < 0:
+                continue
+            
+            # Filter text contain special characters
+            if 'é' in [lb[i].lower() for i in range(len(lb))] or 'ñ' in [lb[i].lower() for i in range(len(lb))] or '£' in [lb[i].lower() for i in range(len(lb))]:
+                continue
+            
+            # Filter out if text is too light or too dark
+            if np.mean(cropped_text) < 30 or np.mean(cropped_text) > 230:
+                cv2.imwrite(f'unvalid_images\\unvalid_image{unvalid_bboxes}_{lb}.jpg', cropped_text)
+                unvalid_bboxes += 1
+                continue
+            
+            # Filter out if image is too small
+            if width < 10 or height < 10:
+                cv2.imwrite(f'unvalid_images\\unvalid_image{unvalid_bboxes}_{lb}.jpg', cropped_text)
+                unvalid_bboxes += 1
+                continue
+            
+            new_image_path = os.path.join(save_dir, f'cropped_image{bboxes_idx}.jpg')
+            cv2.imwrite(new_image_path, cropped_text)
+            new_label = new_image_path + '\t' + lb
+            new_labels.append(new_label)
+            bboxes_idx += 1
+            
+        # Write labels into a text file
+        with open(os.path.join(save_dir, 'labels.txt'), "w") as f:
+            for new_label in new_labels:
+                f.write(f'{new_label}\n')
+                
+
+def build_vocab(root_dir):
+    img_paths = []
+    labels = []
+    
+    # Read labels from text file
+    with open(os.path.join(root_dir, 'ocr_dataset', 'labels.txt'), "r") as f:
+        for label in f:
+            labels.append(label.strip().split("\t")[1])
+            img_paths.append(label.strip().split("\t")[0])
+            
+    # build the vocab
+    vocab = set()
+    for label in labels:
+        for i in range(len(label)):
+            vocab.add(label[i])
+            
+    # "blank" character
+    vocab = list(sorted(vocab))
+    vocab = "".join(vocab)
+    blank_char = '@'
+    vocab = vocab + 'z'
+    vocab = vocab + blank_char
+    
+    # build a dictionary convert from vocab to idx and idx to vocab
+    char_to_idx = {
+        char: idx + 1 for idx, char in enumerate(vocab)
+    }
+    idx_to_char = {
+        idx: char for char, idx in char_to_idx.items()
+    }
+    
+    return char_to_idx, idx_to_char
+    
+def get_imagepaths_and_labels(root_path):
+    img_paths = []
+    labels = []
+    
+    # Read labels from text file
+    with open(os.path.join(root_path, 'ocr_dataset', 'labels.txt'), "r") as f:
+        for label in f:
+            labels.append(label.strip().split("\t")[1])
+            img_paths.append(label.strip().split("\t")[0])
+    
+    return img_paths, labels
+
+def encode(label, char_to_idx, labels):
+    max_length_label = np.max([len(lb) for lb in labels])
+    
+    # encode label
+    encoded_label = torch.tensor(
+                        [char_to_idx[char.lower()] for char in label],
+                        dtype=torch.int32
+                    )
+    label_len = len(encoded_label)
+    length = torch.tensor(
+                label_len,
+                dtype=torch.int32
+            )
+    padded_label = F.pad(
+                        encoded_label,
+                        (0, max_length_label-label_len),
+                        value=0
+                    )
+    return padded_label, length
+
+def decode(encoded_label, idx_to_char, char_to_idx, blank_char='@'):
+    label = []
+    encoded_label = encoded_label.detach().numpy()
+    for i in range(len(encoded_label)):
+        if encoded_label[i] == 0:
+            break
+        elif (i == 0 or encoded_label[i] != encoded_label[i-1]) and encoded_label[i] != char_to_idx[blank_char]:
+            label.append(idx_to_char[encoded_label[i]])
+                
+    label = "".join(label)
+    return label
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--path", type=str, default=os.getcwd(), help="Path to the root directory")
+    args = parser.parse_args()
+    
+    root_path = os.path.join(args.path, 'Dataset')
+
+    image_paths, image_sizes, bboxes, image_labels = extract_data_from_xml(root_path)
+    save_dir = 'Dataset/ocr_dataset'
+    split_bboxes_from_image(image_paths, image_labels, bboxes, save_dir)
+    char_to_idx, idx_to_char = build_vocab(root_path)
+
+if __name__ == '__main__':
+    main()

src/Text_Recognization/text_recognization.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import torchvision
+
+from torchvision.models import resnet101
+
+class BackBone(nn.Module):
+    def __init__(self, num_unfreeze_layers=3):
+        super(BackBone, self).__init__()
+        model = resnet101(weights='IMAGENET1K_V2', progress=True)
+        feature_maps = list(model.children())[:8]
+        
+        # Adding an AdaptiveAvgPooling (batch_size, 2048, 8, 8) -> (batch_size, 2048, 1, 8)
+        feature_maps.append(nn.AdaptiveAvgPool2d((1, None)))
+        self.backbone = nn.Sequential(*feature_maps)
+        
+        for layer in list(self.backbone.parameters())[-(num_unfreeze_layers+1):]:
+            layer.requires_grad = True
+    
+    def forward(self, image):
+        return self.backbone(image)
+
+class CRNN(nn.Module):
+    def __init__(self, vocab_size, hidden_size, n_layers, dropout=0.2, num_unfreeze_layers=3):
+        super(CRNN, self).__init__()
+        self.backbone = BackBone(num_unfreeze_layers=num_unfreeze_layers)
+        
+        self.mapSeq = nn.Sequential(
+            nn.Linear(2048, 512),
+            nn.ReLU(),
+            nn.Dropout(p=dropout)
+        )
+        
+        self.gru = nn.GRU(
+            input_size=512,
+            hidden_size=hidden_size,
+            num_layers=n_layers,
+            bidirectional=True,
+            batch_first=True,
+            dropout=dropout if n_layers > 1 else 0
+        )
+        
+        self.layer_norm = nn.LayerNorm(hidden_size * 2)
+        
+        # Dense layers
+        self.out = nn.Sequential(
+            nn.Linear(hidden_size * 2, vocab_size),
+            nn.LogSoftmax(dim=2)
+        )
+    
+    def forward(self, x):
+        x = self.backbone(x)
+        # (batch_size, 2048, 1, 8) -> (batch_size, 8, 2048, 1)
+        x = x.permute(0, 3, 1, 2)
+        # flatten -> (batch_size, 8, 2048)
+        x = x.view(x.size(0), x.size(1), -1)
+        x = self.mapSeq(x)
+        x, _ = self.gru(x)
+        x = self.layer_norm(x)
+        x = self.out(x)
+        # (batch_size, 8, vocab_size) -> (8, batch_size, vocab_size)
+        x = x.permute(1, 0, 2)
+        
+        return x

src/Text_Recognization/trainer.py

@@ -0,0 +1,162 @@
+import json
+import sys
+import os
+import argparse
+import torch
+import torch.nn as nn
+import torchvision
+from tqdm import tqdm
+
+sys.path.append(os.getcwd())
+
+from src.Text_Recognization.text_recognization import *
+from src.Text_Recognization.prepare_dataset import *
+from src.Text_Recognization.dataloader import *
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+def load_json_config(config_path):
+    with open(config_path, "r") as f:
+        config = json.load(f)
+        
+    return config
+
+def evaluate(model, dataloader, criterion, device):
+    model.eval()
+    losses = []
+    
+    with torch.no_grad():
+        for images, labels, labels_len in dataloader:
+            images = images.to(device)
+            labels = labels.to(device)
+            
+            outputs = model(images)
+            logits_lens = torch.full(
+                size=(outputs.size(1), ),
+                fill_value=outputs.size(0),
+                dtype=torch.long
+            ).to(device)
+            
+            loss = criterion(outputs, labels, logits_lens, labels_len)
+            losses.append(loss.item())
+            
+        eval_loss = sum(losses) / len(losses)
+        return eval_loss
+        
+
+def training_loop(model, train_loader, val_loader, learning_rate, epochs, optimizer, criterion, scheduler, device):
+    model.to(device)
+    
+    train_losses = []
+    val_losses = []
+    
+    for epoch in range(epochs):
+        model.train()
+        
+        batch_losses = []
+        for images, labels, labels_len in tqdm(train_loader):
+            images = images.to(device)
+            labels = labels.to(device)
+            
+            optimizer.zero_grad()
+            outputs = model(images)
+            
+            logits_lens = torch.full(
+                size=(outputs.size(1), ),
+                fill_value=outputs.size(0),
+                dtype=torch.long
+            ).to(device)
+            
+            loss = criterion(outputs, labels, logits_lens, labels_len)
+            
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+            optimizer.step()
+            
+            batch_losses.append(loss.item())
+        
+        train_loss = sum(batch_losses) / len(batch_losses)
+        train_losses.append(train_loss)
+        
+        val_loss = evaluate(model, val_loader, criterion, device)
+        val_losses.append(val_loss)
+        
+        print(f"epoch: {epoch+1}/{epochs}\ttrain_loss:{train_loss}\tval_loss:{val_loss}")
+        
+        scheduler.step()
+        
+    return train_losses, val_losses
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root_path', type=str, default=os.getcwd(), help='Path to the root directory')
+    parser.add_argument('--checkpoints_path', type=str, default=os.path.join(os.getcwd(), 'checkpoints'), help='Path to the checkpoint directory')
+    
+    args = parser.parse_args()
+    config_path = 'src/config.json'
+    dataset_path = os.path.join(args.root_path, 'Dataset')
+    config = load_json_config(config_path)
+
+    # dictionary char and idx
+    char_to_idx, idx_to_char = build_vocab(dataset_path)
+
+    # model
+    model = CRNN(vocab_size=config['vocab_size'], hidden_size=config['CRNN']['hidden_size'], n_layers=config['CRNN']['n_layers'])
+
+    # dataloader
+    train_loader, val_loader, test_loader = get_dataloader()
+
+    # define hyper parammeters
+    criterion = nn.CTCLoss(
+        blank=char_to_idx[config['blank_char']],
+        zero_infinity=True,
+        reduction='mean'
+    )
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=config['CRNN']['learning_rate'],
+        weight_decay=config['CRNN']['weight_decay']
+    )
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer=optimizer,
+        step_size=config['CRNN']['scheduler_step_size'],
+        gamma=0.1
+    )
+    
+    # training loop
+    train_losses, val_losses = training_loop(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        learning_rate=config['CRNN']['learning_rate'],
+        epochs=config['CRNN']['epochs'],
+        optimizer=optimizer,
+        criterion=criterion,
+        scheduler=scheduler,
+        device=device
+    )
+    
+    # save model
+    if not os.path.exists(args.checkpoints_path):
+        os.makedirs(args.checkpoints_path)
+        os.makedirs(os.path.join(args.checkpoints_path, 'losses'))
+    torch.save(model.state_dict(), os.path.join(args.checkpoints_path, 'crnn.pt'))
+    
+    # draw losses
+    fig, axis = plt.subplots(1, 2, figsize=(8, 8))
+    axis[0].plot(train_losses, label='train_loss')
+    axis[0].set_xlabel('Epochs')
+    axis[0].set_ylabel('Loss')
+    axis[0].axis('off')
+    axis[0].legend()
+    
+    axis[1].plot(val_losses, label='val_loss')
+    axis[1].set_xlabel('Epochs')
+    axis[1].set_ylabel('Loss')
+    axis[1].axis('off')
+    axis[1].legend()
+    
+    plt.savefig(os.path.join(args.checkpoints_path, 'losses', 'losses.png'))
+    
+if __name__ == '__main__':
+    main()

src/config.json

@@ -0,0 +1,24 @@
+{
+    "yolov11":
+    {
+        "epochs": 100,
+        "image_size": 640,
+        "cache": true,
+        "patience": 20,
+        "plots": true
+    },
+    "CRNN":
+    {
+        "batch_size": 64,
+        "epochs": 100,
+        "hidden_size": 256,
+        "n_layers": 3,
+        "dropout": 0.2,
+        "unfreeze_layers": 3,
+        "learning_rate": 5e-4,
+        "weight_decay": 1e-5,
+        "scheduler_step_size": 30
+    },
+    "blank_char": "@",
+    "vocab_size": 73
+}

src/predict.py

@@ -0,0 +1,137 @@
+import os
+import sys
+import json
+import cv2
+import argparse
+import matplotlib.pyplot as plt
+
+import ultralytics
+import torch
+import torch.nn as nn
+import torchvision
+from torchvision.models import resnet101
+from torchvision import transforms
+
+sys.path.append(os.getcwd())
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+from ultralytics import YOLO
+from src.Text_Recognization.text_recognization import *
+from src.Text_Recognization.prepare_dataset import *
+
+# config
+def load_json_config(config_path):
+    with open(config_path, "r") as f:
+        config = json.load(f)
+        
+    return config
+
+config = load_json_config('src/config.json')
+
+# char to idx
+char_to_idx, idx_to_char = build_vocab('Dataset')
+
+# text detection model
+text_det_model_path = 'checkpoints/yolov11m.pt'
+yolo = YOLO(text_det_model_path)
+
+# text recognition model
+text_rec_model_path = 'checkpoints/crnn_extend_vocab.pt'
+
+# rcnn model
+rcnn_model = CRNN(vocab_size=74, hidden_size=config['CRNN']['hidden_size'], n_layers=config['CRNN']['n_layers'])
+rcnn_model.load_state_dict(torch.load(text_rec_model_path, weights_only=True, map_location=torch.device('cpu')))
+
+def text_detection(img_path, text_det_model):
+    text_det_results = text_det_model(img_path, verbose=False)[0]
+    
+    bboxes = text_det_results.boxes.xyxy.tolist()
+    classes = text_det_results.boxes.cls.tolist()
+    names = text_det_results.names
+    confs = text_det_results.boxes.conf.tolist()
+    
+    return bboxes, classes, names, confs
+
+def visualize_gt_bboxes_yolo(image_path, gt_location_yolo):
+    image = cv2.imread(image_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    # Convert to original format
+    for data in gt_location_yolo:
+        xmin, ymin, xmax, ymax = data
+        xmin, ymin, xmax, ymax = int(xmin), int(ymin), int(xmax), int(ymax)
+        
+        image = cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color=(255, 0, 0), thickness=2)
+        
+    plt.imshow(image)
+    plt.axis('off')
+    plt.show()
+
+def text_recognization(image, data_transforms, text_reg_model, idx_to_char=idx_to_char, device=device):
+    transformsed_image = data_transforms(image)
+    transformsed_image = transformsed_image.unsqueeze(0).to(device)
+    text_reg_model.to(device)
+    text_reg_model.eval()
+    
+    with torch.no_grad():
+        preds = text_reg_model(transformsed_image)
+        _, idx = torch.max(preds, dim=2)
+        idx = idx.view(-1)
+        text = decode(idx, idx_to_char, char_to_idx)
+        
+    return text, idx
+
+def visualize_detection(image, detections):
+    plt.figure(figsize=(10, 8))
+    
+    for bbox, detected_classes, conf, text, _ in detections:
+        x1, y1, x2, y2 = bbox
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+        
+        image = cv2.rectangle(image, (x1, y1), (x2, y2), color=(255, 0, 0), thickness=2)
+        image = cv2.putText(image, f"{conf:.2f} {text}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+        
+    plt.imshow(image)
+    plt.axis('off')
+    plt.show()
+    return image
+    
+data_transforms = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Resize((100, 400)),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+def prediction(image, text_det_model=yolo, text_reg_model=rcnn_model, idx_to_char=idx_to_char, char_to_idx=char_to_idx, data_transforms=data_transforms, device=device):
+    # detection
+    bboxes, classes, names, confs = text_detection(image, text_det_model)
+    
+    predictions = []
+    for bbox, cls, conf in zip(bboxes, classes, confs):
+        x1, y1, x2, y2 = bbox
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+        detected_text = image[y1:y2, x1:x2]
+        text, encoded_text = text_recognization(detected_text, data_transforms, text_reg_model, idx_to_char, device)
+        predictions.append((bbox, cls, conf, text, encoded_text))
+        print(bbox, cls, conf, text)
+    
+    return predictions
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--image_path', type=str, help='Path to the image')
+    parser.add_argument('--save_path', type=str, default=None, help='Path to save the image')
+    args = parser.parse_args()
+    image_path = args.image_path
+    image = cv2.imread(image_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    detections = prediction(image)
+    image = visualize_detection(image, detections)
+    
+    if args.save_path:
+        print(f"Saving the image to {os.path.join(args.save_path, 'predicted_image.jpg')}")
+        cv2.imwrite(os.path.join(args.save_path, 'predicted_image.jpg'), cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+        
+if __name__ == '__main__':
+    main()