| import torch |
| from torchmetrics import Metric |
| from ocr import OCR |
| import Levenshtein |
|
|
|
|
| class CharacterErrorRate(Metric): |
| def __init__(self, ocr, dist_sync_on_step=False): |
| |
| super().__init__() |
| self.add_state("total_errors", default=torch.tensor(0.0), dist_reduce_fx="sum") |
| self.add_state("total_chars", default=torch.tensor(0.0), dist_reduce_fx="sum") |
| self.ocr = ocr |
|
|
| def update(self, pred_images, target_images): |
| for pred_img, target_img in zip(pred_images, target_images): |
| pred_text = self.ocr.predict(pred_img) |
| target_text = self.ocr.predict(target_img) |
|
|
| dist = Levenshtein.distance(pred_text, target_text) |
| self.total_errors += dist |
| self.total_chars += len(target_text) |
|
|
| def compute(self): |
| if self.total_chars == 0: |
| return torch.tensor(0.0) |
| return self.total_errors / self.total_chars |
|
|
