refactor: update code to lightning>=2.x.x

Author: kwon-evan

lprnet/__pycache__/__init__.cpython-38.pyc

@@ -6,7 +6,7 @@
 import cv2
 from torch.utils.data import Dataset, DataLoader
 from imutils import paths
-import pytorch_lightning as pl
+import lightning as L
 
 from lprnet.utils import encode
 
@@ -24,8 +24,11 @@ def resize_pad(img, size):
         sizeas = (int(w * ash), int(h * ash))
 
     pic1 = cv2.resize(pic1, dsize=sizeas)
-    base_pic[int(size[1] / 2 - sizeas[1] / 2):int(size[1] / 2 + sizeas[1] / 2),
-    int(size[0] / 2 - sizeas[0] / 2):int(size[0] / 2 + sizeas[0] / 2), :] = pic1
+    base_pic[
+        int(size[1] / 2 - sizeas[1] / 2) : int(size[1] / 2 + sizeas[1] / 2),
+        int(size[0] / 2 - sizeas[0] / 2) : int(size[0] / 2 + sizeas[0] / 2),
+        :,
+    ] = pic1
 
     return base_pic
 
@@ -51,20 +54,20 @@ def __init__(self, args, stage, PreprocFun=None):
         self.img_paths = []
         self.img_size = self.args.img_size
 
-        if stage == 'train':
+        if stage == "train":
             self.img_dir = self.args.train_dir
-        elif stage == 'valid':
+        elif stage == "valid":
             self.img_dir = self.args.valid_dir
-        elif stage == 'test':
+        elif stage == "test":
             self.img_dir = self.args.test_dir
-        elif stage == 'predict':
+        elif stage == "predict":
             self.img_dir = self.args.test_dir
         else:
             assert f"No Such Stage. Your input -> {self.stage}"
 
         self.img_paths = [img_path for img_path in paths.list_images(self.img_dir)]
 
-        if stage == 'train':
+        if stage == "train":
             random.shuffle(self.img_paths)
 
         if PreprocFun is not None:
@@ -95,7 +98,7 @@ def __getitem__(self, index):
         return Image, label, len(label)
 
     def transform(self, img):
-        img = img.astype('float32')
+        img = img.astype("float32")
         img -= 127.5
         img *= 0.0078125
         img = np.transpose(img, (2, 0, 1))
@@ -104,13 +107,15 @@ def transform(self, img):
 
     def check(self, label):
         # kor_plate_pattern = re.compile('[가-힣]{0,5}[0-9]{0,3}[가-힣][0-9]{4}')
-        idn_plate_pattern = re.compile('[A-Z]{0,3}[0-9]{0,4}[A-Z]{0,3}')
-        plate_name = idn_plate_pattern.findall(''.join([self.args.chars[c] for c in label]))
+        idn_plate_pattern = re.compile("[A-Z]{0,3}[0-9]{0,4}[A-Z]{0,3}")
+        plate_name = idn_plate_pattern.findall(
+            "".join([self.args.chars[c] for c in label])
+        )
 
         return True if plate_name else False
 
 
-class DataModule(pl.LightningDataModule):
+class DataModule(L.LightningDataModule):
     def __init__(self, args):
         super().__init__()
         self.args = args
@@ -131,29 +136,37 @@ def setup(self, stage: str):
             self.predict = LPRNetDataset(self.args, "predict")
 
     def train_dataloader(self):
-        return DataLoader(self.train,
-                          batch_size=self.args.batch_size,
-                          shuffle=True,
-                          num_workers=4,
-                          collate_fn=collate_fn)
+        return DataLoader(
+            self.train,
+            batch_size=self.args.batch_size,
+            shuffle=True,
+            num_workers=4,
+            collate_fn=collate_fn,
+        )
 
     def val_dataloader(self):
-        return DataLoader(self.val,
-                          batch_size=self.args.batch_size,
-                          shuffle=False,
-                          num_workers=4,
-                          collate_fn=collate_fn)
+        return DataLoader(
+            self.val,
+            batch_size=self.args.batch_size,
+            shuffle=False,
+            num_workers=4,
+            collate_fn=collate_fn,
+        )
 
     def test_dataloader(self):
-        return DataLoader(self.test,
-                          batch_size=self.args.batch_size,
-                          shuffle=False,
-                          num_workers=4,
-                          collate_fn=collate_fn)
+        return DataLoader(
+            self.test,
+            batch_size=self.args.batch_size,
+            shuffle=False,
+            num_workers=4,
+            collate_fn=collate_fn,
+        )
 
     def predict_dataloader(self):
-        return DataLoader(self.predict,
-                          batch_size=self.args.batch_size,
-                          shuffle=False,
-                          num_workers=4,
-                          collate_fn=collate_fn)
+        return DataLoader(
+            self.predict,
+            batch_size=self.args.batch_size,
+            shuffle=False,
+            num_workers=4,
+            collate_fn=collate_fn,
+        )

lprnet/__pycache__/datamodule.cpython-38.pyc

@@ -6,7 +6,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import pytorch_lightning as pl
+import lightning as L
 
 from lprnet.utils import decode, accuracy
 
@@ -33,17 +33,17 @@ def __init__(self):
             nn.Mish(True),
             nn.Conv2d(32, 32, kernel_size=5),
             nn.MaxPool2d(3, stride=3),
-            nn.Mish(True)
+            nn.Mish(True),
         )
         # Regressor for the 3x2 affine matrix
         self.fc_loc = nn.Sequential(
-            nn.Linear(32 * 15 * 6, 32),
-            nn.Mish(True),
-            nn.Linear(32, 3 * 2)
+            nn.Linear(32 * 15 * 6, 32), nn.Mish(True), nn.Linear(32, 3 * 2)
         )
-        # Initialize the weights/bias with identity transformation 
+        # Initialize the weights/bias with identity transformation
         self.fc_loc[2].weight.data.zero_()
-        self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
+        self.fc_loc[2].bias.data.copy_(
+            torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float)
+        )
 
     def forward(self, x):
         xs = self.localization(x)
@@ -99,13 +99,19 @@ def __init__(self, class_num, dropout_rate):
             nn.BatchNorm2d(num_features=256),
             nn.Mish(),
             nn.Dropout(dropout_rate),
-            nn.Conv2d(in_channels=256, out_channels=class_num, kernel_size=(12, 2), stride=1),
+            nn.Conv2d(
+                in_channels=256, out_channels=class_num, kernel_size=(12, 2), stride=1
+            ),
             nn.BatchNorm2d(num_features=class_num),
             nn.Mish(),
         )
         self.container = nn.Sequential(
-            nn.Conv2d(in_channels=256 + class_num + 128 + 64, out_channels=self.class_num, kernel_size=(1, 1),
-                      stride=(1, 1)),
+            nn.Conv2d(
+                in_channels=256 + class_num + 128 + 64,
+                out_channels=self.class_num,
+                kernel_size=(1, 1),
+                stride=(1, 1),
+            ),
         )
 
     def forward(self, x):
@@ -133,12 +139,14 @@ def forward(self, x):
         return logits
 
 
-class LPRNet(pl.LightningModule):
+class LPRNet(L.LightningModule):
     def __init__(self, args: Optional[Namespace] = None):
         super().__init__()
         self.save_hyperparameters(args)
         self.STNet = _STNet()
-        self.LPRNet = _LPRNet(class_num=len(self.hparams.chars), dropout_rate=self.hparams.dropout_rate)
+        self.LPRNet = _LPRNet(
+            class_num=len(self.hparams.chars), dropout_rate=self.hparams.dropout_rate
+        )
 
     def forward(self, x):
         return self.LPRNet(self.STNet(x))
@@ -150,10 +158,17 @@ def training_step(self, batch, batch_idx):
         logits = self(imgs)
         log_probs = logits.permute(2, 0, 1)
         log_probs = log_probs.log_softmax(2).requires_grad_()
-        input_lengths, target_lengths = sparse_tuple_for_ctc(self.hparams.t_length, lengths)
-        loss = F.ctc_loss(log_probs=log_probs, targets=labels,
-                          input_lengths=input_lengths, target_lengths=target_lengths,
-                          blank=len(self.hparams.chars) - 1, reduction='mean')
+        input_lengths, target_lengths = sparse_tuple_for_ctc(
+            self.hparams.t_length, lengths
+        )
+        loss = F.ctc_loss(
+            log_probs=log_probs,
+            targets=labels,
+            input_lengths=input_lengths,
+            target_lengths=target_lengths,
+            blank=len(self.hparams.chars) - 1,
+            reduction="mean",
+        )
         acc = accuracy(logits, labels, lengths, self.hparams.chars)
 
         self.log("train-loss", abs(loss), prog_bar=True, logger=True, sync_dist=True)
@@ -167,10 +182,17 @@ def validation_step(self, batch, batch_idx):
         logits = self(imgs)
         log_probs = logits.permute(2, 0, 1)
         log_probs = log_probs.log_softmax(2).requires_grad_()
-        input_lengths, target_lengths = sparse_tuple_for_ctc(self.hparams.t_length, lengths)
-        loss = F.ctc_loss(log_probs=log_probs, targets=labels,
-                          input_lengths=input_lengths, target_lengths=target_lengths,
-                          blank=len(self.hparams.chars) - 1, reduction='mean')
+        input_lengths, target_lengths = sparse_tuple_for_ctc(
+            self.hparams.t_length, lengths
+        )
+        loss = F.ctc_loss(
+            log_probs=log_probs,
+            targets=labels,
+            input_lengths=input_lengths,
+            target_lengths=target_lengths,
+            blank=len(self.hparams.chars) - 1,
+            reduction="mean",
+        )
         acc = accuracy(logits, labels, lengths, self.hparams.chars)
 
         self.log("val-loss", abs(loss), prog_bar=True, logger=True, sync_dist=True)
@@ -179,14 +201,22 @@ def validation_step(self, batch, batch_idx):
     def test_step(self, batch, batch_idx):
         imgs, labels, lengths = batch
         import time
+
         start = time.time()
         logits = self(imgs)
         log_probs = logits.permute(2, 0, 1)
         log_probs = log_probs.log_softmax(2).requires_grad_()
-        input_lengths, target_lengths = sparse_tuple_for_ctc(self.hparams.t_length, lengths)
-        loss = F.ctc_loss(log_probs=log_probs, targets=labels,
-                          input_lengths=input_lengths, target_lengths=target_lengths,
-                          blank=len(self.hparams.chars) - 1, reduction='mean')
+        input_lengths, target_lengths = sparse_tuple_for_ctc(
+            self.hparams.t_length, lengths
+        )
+        loss = F.ctc_loss(
+            log_probs=log_probs,
+            targets=labels,
+            input_lengths=input_lengths,
+            target_lengths=target_lengths,
+            blank=len(self.hparams.chars) - 1,
+            reduction="mean",
+        )
         acc = accuracy(logits, labels, lengths, self.hparams.chars)
         end = time.time()
 
@@ -204,17 +234,27 @@ def predict_step(self, batch, batch_idx, dataloader_idx: int = 0):
         return predict
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam([{'params': self.STNet.parameters(),
-                                       'weight_decay': self.hparams.weight_decay},
-                                      {'params': self.LPRNet.parameters()}],
-                                     lr=self.hparams.lr)
-        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 10, 2, 0.0001, -1)
-        return {"optimizer": optimizer,
-                "lr_scheduler": {
-                    "scheduler": scheduler,
-                    "interval": "step",
-                    "frequency": 1,
-                    "monitor": "val-loss",
-                    "strict": True,
-                    "name": "lr"
-                }}
+        optimizer = torch.optim.Adam(
+            [
+                {
+                    "params": self.STNet.parameters(),
+                    "weight_decay": self.hparams.weight_decay,
+                },
+                {"params": self.LPRNet.parameters()},
+            ],
+            lr=self.hparams.lr,
+        )
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            optimizer, 10, 2, 0.0001, -1
+        )
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",
+                "frequency": 1,
+                "monitor": "val-loss",
+                "strict": True,
+                "name": "lr",
+            },
+        }

lprnet/__pycache__/lprnet.cpython-38.pyc

@@ -17,7 +17,7 @@ def encode(imgname: str, chars: List[str]):
             label.append(chars_dict[imgname[i:j]])
             i = j
         else:
-            assert 0, f'no such char in {imgname}'
+            assert 0, f"no such char in {imgname}"
 
     return label
 
@@ -32,7 +32,7 @@ def decode(preds, chars):
         for j in range(pred.shape[1]):
             pred_label.append(np.argmax(pred[:, j], axis=0))
         no_repeat_blank_label = list()
-        pre_c = ''
+        pre_c = ""
         for c in pred_label:  # dropout repeated label and blank label
             if (pre_c == c) or (c == len(chars) - 1):
                 if c == len(chars) - 1:
@@ -58,7 +58,7 @@ def accuracy(logits, labels, lengths, chars):
     TP, total = 0, 0
     start = 0
     for i, length in enumerate(lengths):
-        label = labels[start:start + length]
+        label = labels[start : start + length]
         start += length
         if np.array_equal(np.array(pred_labels[i]), label.cpu().numpy()):
             TP += 1
@@ -72,19 +72,20 @@ def tensor2numpy(inp):
     inp = inp.squeeze(0).cpu()
     inp = inp.detach().numpy().transpose((1, 2, 0))
     inp = 127.5 + inp / 0.0078125
-    inp = inp.astype('uint8')
+    inp = inp.astype("uint8")
 
     return inp
 
 
 def numpy2tensor(img: np.ndarray, img_size: Sequence[int]):
     # convert a numpy image to tensor
     import cv2
+
     height, width, _ = img.shape
 
     if height != img_size[1] or width != img_size[0]:
         img = cv2.resize(img, img_size, interpolation=cv2.INTER_CUBIC)
-    img = img.astype('float32')
+    img = img.astype("float32")
     img -= 127.5
     img *= 0.0078125
     img = np.transpose(img, (2, 0, 1))

lprnet/__pycache__/utils.cpython-38.pyc

@@ -7,15 +7,16 @@
 import warnings
 import yaml
 import torch
-from pytorch_lightning import Trainer
+
+import lightning as L
 
 from lprnet import LPRNet, DataModule
 
 warnings.filterwarnings("ignore")
 
 
-if __name__ == '__main__':
-    with open('config/idn_config.yaml') as f:
+if __name__ == "__main__":
+    with open("config/idn_config.yaml") as f:
         args = Namespace(**yaml.load(f, Loader=yaml.FullLoader))
 
     load_model_start = time.time()
@@ -29,7 +30,7 @@
 
     dm = DataModule(args)
 
-    trainer = Trainer(
+    trainer = L.Trainer(
         accelerator="auto",
         precision=16,
         devices=torch.cuda.device_count(),