Update ocr training

ria-com · Jan 6, 2025 · 90f1857 · 90f1857
1 parent 3645625
commit 90f1857
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/nomeroff_net/data_loaders/text_image_generator.py b/nomeroff_net/data_loaders/text_image_generator.py
@@ -92,12 +92,12 @@ def generate_cache_x_in_path(self, img_path: str, cache_dirpath: str, newsize: T
         if newsize is None:
             newsize = (self.img_w, self.img_h)
         img = Image.open(img_path).convert('RGB')
-        img = img.resize(newsize)
         if self.with_aug:
             from nomeroff_net.tools.augmentations import aug
             img = np.array(img)
             imgs = aug([img])
             img = Image.fromarray(imgs[0])
+        img = img.resize(newsize)
         x = self.transform(img)
         torch.save(x, x_path)
         return x_path

diff --git a/nomeroff_net/pipes/number_plate_text_readers/base/ocr.py b/nomeroff_net/pipes/number_plate_text_readers/base/ocr.py
@@ -182,6 +182,7 @@ def train(self,
             pl.seed_everything(seed)
         if self.model is None:
             self.create_model()
+        self.model.train()
         checkpoint_callback = ModelCheckpoint(dirpath=log_dir, monitor='val_loss')
         lr_monitor = LearningRateMonitor(logging_interval='step')
         if self.gpus: