Merge branch 'v3.4'

# Conflicts:
#	README.md
#	nomeroff_net/__init__.py

History.md

@@ -1,3 +1,10 @@
+3.4.0 / 2022-11-21
+==================
+  **updates**
+  * Update ocr models
+  * Added support yolov8 models for numberpalte detection
+  * Update option classification models with efficientnet_v2_s backbone 
+
 3.3.0 / 2022-11-21
 ==================
   **updates**

README.md

@@ -1,6 +1,6 @@
 ![Nomeroff Net. Automatic numberplate recognition system](./public/images/nomeroff_net.svg)
 
-Nomeroff Net. Automatic numberplate recognition system. Version 3.3
+Nomeroff Net. Automatic numberplate recognition system. Version 3.4
 <br /><br />
 <blockquote style="border-left-color: #ff0000">
 Now there is a war going on in our country, Russian soldiers are shooting at civilians in Ukraine. Enemy aviation launches rockets and drops bombs on residential quarters.

nomeroff_net/__init__.py

@@ -9,4 +9,4 @@
 from nomeroff_net.pipelines import pipeline
 
 
-__version__ = "3.3.0"
+__version__ = "3.4.0"

nomeroff_net/nnmodels/numberplate_options_model.py

@@ -6,7 +6,7 @@
 import torch.nn as nn
 from torch.nn import functional
 from .numberplate_classification_model import ClassificationNet
-from torchvision.models import resnet18
+from torchvision.models import efficientnet_v2_s
 from nomeroff_net.tools.errors import NPOptionsNetError
 import contextlib
 from nomeroff_net.tools.mcm import get_device_torch
@@ -17,36 +17,63 @@ def dummy_context_mgr():
     yield None
 
 
+class DoubleLinear(torch.nn.Module):
+    def __init__(self, linear1, linear2):
+        super(DoubleLinear, self).__init__()
+        self.linear1 = linear1
+        self.linear2 = linear2
+
+    def forward(self, input):
+        return self.linear1(input), self.linear2(input)
+
+
 class NPOptionsNet(ClassificationNet):
     def __init__(self,
                  region_output_size: int,
                  count_line_output_size: int,
                  batch_size: int = 1,
-                 learning_rate: float = 0.005,
+                 learning_rate: float = 0.001,
                  train_regions=True,
-                 train_count_lines=True):
+                 train_count_lines=True,
+                 backbone=None):
         super(NPOptionsNet, self).__init__() 
         self.batch_size = batch_size
         self.learning_rate = learning_rate
 
         self.train_regions = train_regions
         self.train_count_lines = train_count_lines
 
-        resnet = resnet18(pretrained=True)
-        modules = list(resnet.children())[:-3]
-        self.resnet = nn.Sequential(*modules)
+        if backbone is None:
+            backbone = efficientnet_v2_s
+        self.model = backbone()
 
-        self.dropout_reg = nn.Dropout(0.2)
-        self.fc1_reg = nn.Linear(256 * 4 * 19, 512)
-        self.fc2_reg = nn.Linear(512, 256)
-        self.batch_norm_reg = nn.BatchNorm1d(512)
-        self.fc3_reg = nn.Linear(256, region_output_size)
-
-        self.dropout_line = nn.Dropout(0.2)
-        self.fc1_line = nn.Linear(256 * 4 * 19, 512)
-        self.fc2_line = nn.Linear(512, 256)
-        self.batch_norm_line = nn.BatchNorm1d(512)
-        self.fc3_line = nn.Linear(256, count_line_output_size)
+        if 'efficientnet' in str(backbone):
+            in_features = self.model.classifier[1].in_features
+        else:
+            raise NotImplementedError(backbone)
+
+        linear_region = nn.Sequential(
+            nn.Dropout(p=0.2, inplace=False),
+            nn.Linear(in_features=in_features,
+                      out_features=region_output_size,
+                      bias=True)
+        )
+        if not self.train_regions:
+            for name, param in linear_region.named_parameters():
+                param.requires_grad = False
+        linear_line = nn.Sequential(
+            nn.Dropout(p=0.2, inplace=False),
+            nn.Linear(in_features=in_features,
+                      out_features=count_line_output_size,
+                      bias=True)
+        )
+        if not self.train_count_lines:
+            for name, param in linear_line.named_parameters():
+                param.requires_grad = False
+        if 'efficientnet' in str(backbone):
+            self.model.classifier = DoubleLinear(linear_region, linear_line)
+        else:
+            raise NotImplementedError(backbone)
 
     def training_step(self, batch, batch_idx):
         loss, acc, acc_reg, acc_line = self.step(batch)
@@ -104,25 +131,10 @@ def test_step(self, batch, batch_idx):
         }
 
     def forward(self, x):
-        x = self.resnet(x)
-
-        with dummy_context_mgr() if self.train_regions else torch.no_grad():
-            x1 = x.reshape(x.size(0), -1)
-            x1 = self.dropout_reg(x1)
-            x1 = functional.relu(self.fc1_reg(x1))
-            if self.batch_size > 1:
-                x1 = self.batch_norm_reg(x1)
-            x1 = functional.relu(self.fc2_reg(x1))
-            x1 = functional.softmax(self.fc3_reg(x1))
-
-        with dummy_context_mgr() if self.train_count_lines else torch.no_grad():
-            x2 = x.reshape(x.size(0), -1)
-            x2 = self.dropout_line(x2)
-            x2 = functional.relu(self.fc1_line(x2))
-            if self.batch_size > 1:
-                x2 = self.batch_norm_line(x2)
-            x2 = functional.relu(self.fc2_line(x2))
-            x2 = functional.softmax(self.fc3_line(x2))
+
+        x1, x2 = self.model(x)
+        x1 = functional.softmax(x1)
+        x2 = functional.softmax(x2)
 
         return x1, x2
 
@@ -151,12 +163,19 @@ def step(self, batch):
         return loss, acc, acc_reg, acc_line
 
     def configure_optimizers(self):
-        optimizer = torch.optim.ASGD(self.parameters(),
-                                     lr=self.learning_rate,
-                                     lambd=0.0001,
-                                     alpha=0.75,
-                                     t0=1000000.0,
-                                     weight_decay=0)
+        optimizer = torch.optim.SGD(self.parameters(), lr=self.learning_rate, momentum=0.9)
+
+        # # Try this
+        # from lion_pytorch import Lion
+        # optimizer = Lion(self.parameters(), lr=self.learning_rate)
+
+        # # Old optimizer
+        # optimizer = torch.optim.ASGD(self.parameters(),
+        #                              lr=self.learning_rate,
+        #                              lambd=0.0001,
+        #                              alpha=0.75,
+        #                              t0=1000000.0,
+        #                              weight_decay=0)
         return optimizer
 
 

nomeroff_net/pipelines/multiline_number_plate_detection_and_reading.py

@@ -43,7 +43,7 @@ def forward(self, inputs: Any, **forward_parameters: Dict) -> Any:
          count_lines, confidences, predicted,
          zones, image_ids,
          images_bboxs, images,
-         images_points, images_mline_boxes) = self.forward_detection_np(inputs, **forward_parameters)
+         images_points, images_mline_boxes, preprocessed_np) = self.forward_detection_np(inputs, **forward_parameters)
         zones = convert_multiline_images_to_one_line(
             image_ids,
             images,
@@ -56,4 +56,4 @@ def forward(self, inputs: Any, **forward_parameters: Dict) -> Any:
                                            count_lines, confidences,
                                            zones, image_ids,
                                            images_bboxs, images,
-                                           images_points, **forward_parameters)
+                                           images_points, preprocessed_np, **forward_parameters)

nomeroff_net/pipelines/number_plate_classification.py

@@ -35,10 +35,14 @@ def preprocess(self, inputs: Any, **preprocess_parameters: Dict) -> Any:
     @no_grad()
     def forward(self, inputs: Any, **forward_parameters: Dict) -> Any:
         model_output = self.detector.forward(inputs)
-        return unzip([p.cpu().numpy() for p in model_output])
+        model_output = [p.cpu().numpy() for p in model_output]
+        model_output = [*model_output, inputs]
+        return unzip(model_output)
 
     def postprocess(self, inputs: Any, **postprocess_parameters: Dict) -> Any:
-        confidences, region_ids, count_lines = self.detector.unzip_predicted(unzip(inputs))
+        unziped_inputs = unzip(inputs)
+        processed_np = [np for np in unziped_inputs[2]]
+        confidences, region_ids, count_lines = self.detector.unzip_predicted(unziped_inputs)
         count_lines = self.detector.custom_count_lines_id_to_all_count_lines(count_lines)
         region_names = self.detector.get_region_labels(region_ids)
-        return unzip([region_ids, region_names, count_lines, confidences, inputs])
+        return unzip([region_ids, region_names, count_lines, confidences, inputs, processed_np])

nomeroff_net/pipelines/number_plate_detection_and_reading.py

@@ -27,13 +27,16 @@ def __init__(self,
                  default_label: str = "eu_ua_2015",
                  default_lines_count: int = 1,
                  number_plate_localization_class: Pipeline = DefaultNumberPlateLocalization,
+                 number_plate_localization_detector=None,
                  **kwargs):
         self.default_label = default_label
         self.default_lines_count = default_lines_count
         self.number_plate_localization = number_plate_localization_class(
             "number_plate_localization",
             image_loader=None,
-            path_to_model=path_to_model)
+            path_to_model=path_to_model,
+            detector=number_plate_localization_detector
+        )
         self.number_plate_key_points_detection = NumberPlateKeyPointsDetection(
             "number_plate_key_points_detection",
             image_loader=None,
@@ -81,22 +84,24 @@ def forward_detection_np(self, inputs: Any, **forward_parameters: Dict):
             count_lines = [self.default_lines_count for _ in zones]
             confidences = [-1 for _ in zones]
             predicted = [-1 for _ in zones]
+            preprocessed_np = [None for _ in zones]
         else:
             (region_ids, region_names, count_lines,
-             confidences, predicted) = unzip(self.number_plate_classification(zones, **forward_parameters))
+             confidences, predicted, preprocessed_np) = unzip(self.number_plate_classification(zones,
+                                                                                               **forward_parameters))
         return (region_ids, region_names, count_lines, confidences,
                 predicted, zones, image_ids, images_bboxs, images,
-                images_points, images_mline_boxes)
+                images_points, images_mline_boxes, preprocessed_np)
 
     def forward_recognition_np(self, region_ids, region_names,
                                count_lines, confidences,
                                zones, image_ids,
                                images_bboxs, images,
-                               images_points, **forward_parameters):
+                               images_points, preprocessed_np, **forward_parameters):
         number_plate_text_reading_res = unzip(
             self.number_plate_text_reading(unzip([zones,
                                                   region_names,
-                                                  count_lines]), **forward_parameters))
+                                                  count_lines, preprocessed_np]), **forward_parameters))
         if len(number_plate_text_reading_res):
             texts, _ = number_plate_text_reading_res
         else:
@@ -116,12 +121,12 @@ def forward(self, inputs: Any, **forward_parameters: Dict) -> Any:
          count_lines, confidences, predicted,
          zones, image_ids,
          images_bboxs, images,
-         images_points, images_mline_boxes) = self.forward_detection_np(inputs, **forward_parameters)
+         images_points, images_mline_boxes, preprocessed_np) = self.forward_detection_np(inputs, **forward_parameters)
         return self.forward_recognition_np(region_ids, region_names,
                                            count_lines, confidences,
                                            zones, image_ids,
                                            images_bboxs, images,
-                                           images_points, **forward_parameters)
+                                           images_points, preprocessed_np, **forward_parameters)
 
     @empty_method
     def postprocess(self, inputs: Any, **postprocess_parameters: Dict) -> Any:

nomeroff_net/pipelines/number_plate_localization.py

@@ -3,7 +3,8 @@
 from nomeroff_net.image_loaders import BaseImageLoader
 from nomeroff_net.pipelines.base import Pipeline
 from nomeroff_net.tools import unzip
-from nomeroff_net.pipes.number_plate_localizators.yolo_v5_detector import Detector
+#from nomeroff_net.pipes.number_plate_localizators.yolo_v5_detector import Detector
+from nomeroff_net.pipes.number_plate_localizators.yolo_v8_detector import Detector
 
 
 class NumberPlateLocalization(Pipeline):
@@ -15,9 +16,12 @@ def __init__(self,
                  task,
                  image_loader: Optional[Union[str, BaseImageLoader]],
                  path_to_model="latest",
+                 detector=None,
                  **kwargs):
         super().__init__(task, image_loader, **kwargs)
-        self.detector = Detector()
+        if detector is None:
+            detector = Detector
+        self.detector = detector()
         self.detector.load(path_to_model)
 
     def sanitize_parameters(self, img_size=None, stride=None, min_accuracy=None, **kwargs):

nomeroff_net/pipelines/number_plate_text_reading.py

@@ -19,15 +19,15 @@
         "for_regions": ["eu", "xx_transit", "xx_unknown"],
         "model_path": "latest"
     },
-    "ru_shufflenet_v2_x2_0": {
+    "ru": {
         "for_regions": ["ru", "eu_ua_ordlo_lpr", "eu_ua_ordlo_dpr"],
         "model_path": "latest"
     },
-    "kz_shufflenet_v2_x2_0": {
+    "kz": {
         "for_regions": ["kz"],
         "model_path": "latest"
     },
-    "kg_shufflenet_v2_x2_0": {
+    "kg": {  # "kg_shufflenet_v2_x2_0"
         "for_regions": ["kg"],
         "model_path": "latest"
     },
@@ -39,7 +39,7 @@
         "for_regions": ["su"],
         "model_path": "latest"
     },
-    "am_shufflenet_v2_x2_0": {
+    "am": {
         "for_regions": ["am"],
         "model_path": "latest"
     },
@@ -62,11 +62,13 @@ def __init__(self,
                  default_label: str = "eu_ua_2015",
                  default_lines_count: int = 1,
                  class_detector=TextDetector,
+                 need_preprocess=False,
                  **kwargs):
         if presets is None:
             presets = DEFAULT_PRESETS
         super().__init__(task, image_loader, **kwargs)
         self.detector = class_detector(presets, default_label, default_lines_count)
+        self.need_preprocess = need_preprocess
 
     def sanitize_parameters(self, **kwargs):
         return {}, {}, {}
@@ -75,14 +77,17 @@ def __call__(self, images: Any, **kwargs):
         return super().__call__(images, **kwargs)
 
     def preprocess(self, inputs: Any, **preprocess_parameters: Dict) -> Any:
-        images, labels, lines = unzip(inputs)
+        images, labels, lines, preprocessed_np = unzip(inputs)
         images = [self.image_loader.load(item) for item in images]
-        return unzip([images, labels, lines])
+        return unzip([images, labels, lines, preprocessed_np])
 
     @no_grad()
     def forward(self, inputs: Any, **forward_parameters: Dict) -> Any:
-        images, labels, lines = unzip(inputs)
-        model_inputs = self.detector.preprocess(images, labels, lines)
+        images, labels, lines, preprocessed_np = unzip(inputs)
+        if self.need_preprocess or all([np is None for np in preprocessed_np]):
+            model_inputs = self.detector.preprocess(images, labels, lines)
+        else:
+            model_inputs = self.detector.preprocess(preprocessed_np, labels, lines, need_preprocess=False)
         model_outputs = self.detector.forward(model_inputs)
         model_outputs = self.detector.postprocess(model_outputs)
         return unzip([images, model_outputs, labels])

nomeroff_net/pipes/number_plate_classificators/options_detector.py

@@ -75,8 +75,8 @@ def __init__(self, options: Dict = None) -> None:
             options = dict()
 
         # input
-        self.height = 64
-        self.width = 295
+        self.height = 50
+        self.width = 200
         self.color_channels = 3
 
         # outputs 1
@@ -398,9 +398,8 @@ def unzip_predicted(predicted):
             confidences.append([region_confidence, count_lines_confidence])
         return confidences, region_ids, count_lines
 
-    @staticmethod
-    def preprocess(images):
-        x = [normalize_img(img) for img in images]
+    def preprocess(self, images):
+        x = [normalize_img(img, height=self.height, width=self.width) for img in images]
         x = np.moveaxis(np.array(x), 3, 1)
         return x
 
@@ -415,7 +414,7 @@ def predict_with_confidence(self, imgs: List[np.ndarray or List]) -> Tuple:
         """
         Predict options(region, count lines) with confidence by numberplate images
         """
-        xs = [normalize_img(img) for img in imgs]
+        xs = [normalize_img(img, height=self.height, width=self.width) for img in imgs]
         if not bool(xs):
             return [], [], [], []
         predicted = self._predict(xs)