classify_and_transribe.py

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image
import os
import argparse
from pathlib import Path


classes_to_chars = {
      0: "Ι",
      1: "Θ",
      2: "Κ",
      3: "Λ",
      4: "Ϛ",
      5: "Υ",
      6: "Α",
      7: "Ε",
      8: "Η",
      9: "Τ",
      10: "Ρ",
      11: "Ο",
      12: "",
      13: "Ν",
      14: "Ω",
      15: "Γ",
      16: "Π",
      17: "Β",
      18: "Φ",
      19: "ΟΥ",
      20: "+",
      21: "Μ",
      22: "Δ",
      23: "Χ",
      24: "S"
}



# Image transformation
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

def load_model(modelf):

    # Load Resnet18 model
    model = models.resnet18()
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 25)
    model.load_state_dict(torch.load(modelf, map_location=torch.device('cpu')))
    model.eval()

    return model


# Inference
def predict_image(model, image_path):
    image = Image.open(image_path)
    image = transform(image).unsqueeze(0)

    with torch.no_grad():
        outputs = model(image)
        _, predicted = outputs.max(1)
        return predicted.item()



list_chars = ['Ι', 'Θ', 'Κ', 'Λ', 'C moon-shaped sigma', 'V = Y',
              'Α', 'Ε', 'Η', 'Τ', 'Ρ = rho', 'Ο', 'bg', 'Ν', 'ω',
              'Γ', 'Π', 'R = βῆτα', 'Φ', 'ligature OU', 'Croisette',
              'Μ', 'Δ', 'Χ', 'S = καί']




def classify(model, boxesf, cropsf):

    # Open boxes file
    boxes = []
    with open(boxesf, 'r') as f:
        for line in f:
            fields = line.split()
            if len(fields) > 1:
                boxes.append( [ int(fields[0]), tuple(float(x) for x in fields[1:])])

    # print(f"Boxes: {boxes}")


    # Apply model to crops
    classified_boxes = []
    stem = Path(boxesf).stem

    for i, box in enumerate(boxes):


        suffix = "" if i==0 else str(i+1)

        cropfile = os.path.join(cropsf, f"{stem}{suffix}.jpg")

        if not os.path.isfile(cropfile):
            print("Error trying to find crop file: ", cropfile)
            return

        prediction = predict_image(model, cropfile)
        classified_boxes.append([prediction, *box[1]])
        print(f"Classified box {i} to {prediction}")

    # print(f"Classified Boxes: {classified_boxes}")


    # Update box file with results
    with open(boxesf, 'w') as f:
        for box in classified_boxes:
            f.write(" ".join(str(x) for x in box))
            f.write("\n")



    return classified_boxes



def transcript(boxes, keysf, outfile):

    keys = []
    with open(keysf, "r") as f:
        for line in f:
            keys.append( [ int(x) for x in line.split()] )


    transcription = []

    for line in keys:

        transcription.append( "".join( classes_to_chars[  boxes[k][0] ]    for k in line ) )

    transtr = "\n".join(transcription)

    print("transcription: ")
    print(transtr)

    with open(outfile, "w") as fout:
        fout.write(transtr)

    return transtr



if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('--model', required=True, help="trained model to use")
    parser.add_argument('--boxes', required=True, help="txt file containing detected boxes")
    parser.add_argument('--crops', required=True, help="path containing the cropped images of corresponding the boxes")

    parser.add_argument('--keys', help="ordered set of box indices for transcription")

    parser.add_argument('--out', help="ouput directory")

    args = parser.parse_args()

    if not args.out:
        args.out = str( Path(args.boxes).parent )

    os.makedirs(args.out, exist_ok=True)

    outfile = os.path.join(args.out, Path(args.boxes).stem + ".txt" )

    if os.path.isfile(args.model) and os.path.isfile(args.boxes) and os.path.isdir(args.crops):


        model = load_model(args.model)

        classified_boxes = classify(model, args.boxes, args.crops)

        if args.keys and os.path.isfile(args.keys):
            transcript(classified_boxes, args.keys, outfile)
    else:
        parser.print_usage()
        exit(1)