Improve style.

Author: Adibvafa

openpmcvl/granular/models/network.py

@@ -1,5 +1,5 @@
 import torch
-import torch.nn as nn
+from torch import nn
 from torch.utils.model_zoo import load_url as load_state_dict_from_url
 
 

openpmcvl/granular/models/process.py

@@ -1,8 +1,8 @@
 from __future__ import division
 
-import torch
 import cv2
 import numpy as np
+import torch
 
 
 def label2yolobox(labels, info_img, maxsize, lrflip):
@@ -22,7 +22,8 @@ class (float): class index.
         maxsize (int): target image size after pre-processing
         lrflip (bool): horizontal flip flag
 
-    Returns:
+    Returns
+    -------
         labels:label data whose size is :math:`(N, 5)`.
             Each label consists of [class, xc, yc, w, h] where
                 class (float): class index.
@@ -77,7 +78,9 @@ def nms(bbox, thresh, score=None, limit=None):
         limit (int): The upper bound of the number of the output bounding
             boxes. If it is not specified, this method selects as many
             bounding boxes as possible.
-    Returns:
+
+    Returns
+    -------
         array:
         An array with indices of bounding boxes that are selected. \
         They are sorted by the scores of bounding boxes in descending \
@@ -87,7 +90,6 @@ def nms(bbox, thresh, score=None, limit=None):
 
     from: https://github.com/chainer/chainercv
     """
-
     if len(bbox) == 0:
         return np.zeros((0,), dtype=np.int32)
 
@@ -135,7 +137,8 @@ def postprocess(prediction, dtype, conf_thre=0.7, nms_thre=0.45):
         nms_thre (float):
             IoU threshold of non-max suppression ranging from 0 to 1.
 
-    Returns:
+    Returns
+    -------
         output (list of torch tensor):
 
     """
@@ -203,7 +206,8 @@ def preprocess(img, imgsize, jitter, random_placing=False):
         jitter (float): amplitude of jitter for resizing
         random_placing (bool): if True, place the image at random position
 
-    Returns:
+    Returns
+    -------
         img (numpy.ndarray): input image whose shape is :math:`(C, imgsize, imgsize)`.
             Values range from 0 to 1.
         info_img : tuple of h, w, nh, nw, dx, dy.

openpmcvl/granular/models/subfigure_detector.py

@@ -1,11 +1,12 @@
+import math
+
 import torch
+from einops import repeat
 from pytorch_pretrained_bert.modeling import BertModel
 from torch import nn
 from torchvision import models
-import math
 
 from openpmcvl.granular.models.transformer_module import *
-from einops import repeat
 
 
 class FigCap_Former(nn.Module):
@@ -116,7 +117,8 @@ def forward(self, images, texts):
             images (compound figure): shape (bs, c, h, w)
             texts (caption tokens): shape (bs, max_length_in_this_batch)
 
-        Returns:
+        Returns
+        -------
             output_det_class: tensor (bs, query_num, 1), 0~1 indicate subfigure or no-subfigure
             output_box: tensor (bs, query_num, 4), prediction of [cx, cy, w, h]
             similarity: tensor (bs, query_num, caption_length), 0~1 indicate belong or not belong to the subfigure

openpmcvl/granular/models/subfigure_ocr.py

@@ -1,16 +1,17 @@
 import os
-import yaml
+
+import cv2
+import numpy as np
 import torch
+import torch.nn.functional as F
+import yaml
+from PIL import Image
 from skimage import io
-import numpy as np
-import cv2
 from torch.autograd import Variable
-from PIL import Image
-import torch.nn.functional as F
 
-from openpmcvl.granular.models.yolov3 import YOLOv3
 from openpmcvl.granular.models.network import resnet152
-from openpmcvl.granular.models.process import preprocess, postprocess, yolobox2label
+from openpmcvl.granular.models.process import postprocess, preprocess, yolobox2label
+from openpmcvl.granular.models.yolov3 import YOLOv3
 
 
 class classifier:
@@ -44,7 +45,7 @@ def __init__(self):
         self.text_recognition_model.eval()
 
     def load_model_from_checkpoint(self, model, model_name):
-        """load checkpoint weights into model"""
+        """Load checkpoint weights into model"""
         checkpoints_path = os.path.join(self.current_dir, "..", "checkpoints")
         checkpoint = os.path.join(checkpoints_path, model_name)
         model.load_state_dict(torch.load(checkpoint))
@@ -61,7 +62,6 @@ def detect_subfigure_boundaries(self, figure_path):
             subfigure_info (list of lists): Each inner list is
                 x1, y1, x2, y2, confidence
         """
-
         ## Preprocess the figure for the models
         img = io.imread(figure_path)
         if len(np.shape(img)) == 2:

openpmcvl/granular/models/transformer_module.py

@@ -1,7 +1,8 @@
+import copy
+
 import torch
-from torch import nn
 import torch.nn.functional as F
-import copy
+from torch import nn
 
 
 class MultiHeadAttention(nn.Module):

openpmcvl/granular/models/yolo_layer.py

@@ -1,7 +1,9 @@
-import torch
-import torch.nn as nn
-import numpy as np
 import warnings
+
+import numpy as np
+import torch
+from torch import nn
+
 from openpmcvl.granular.models.network import resnet152
 from openpmcvl.granular.models.process import preprocess
 
@@ -18,7 +20,9 @@ def bboxes_iou(bboxes_a, bboxes_b, xyxy=True):
         bbox_b (array): An array similar to :obj:`bbox_a`,
             whose shape is :math:`(K, 4)`.
             The dtype should be :obj:`numpy.float32`.
-    Returns:
+
+    Returns
+    -------
         array:
         An array whose shape is :math:`(N, K)`. \
         An element at index :math:`(n, k)` contains IoUs between \
@@ -72,7 +76,6 @@ def __init__(self, config_model, layer_no, in_ch, ignore_thre=0.7):
             in_ch (int): number of input channels.
             ignore_thre (float): threshold of IoU above which objectness training is ignored.
         """
-
         super(YOLOLayer, self).__init__()
         strides = [32, 16, 8]  # fixed
         self.anchors = config_model["ANCHORS"]
@@ -111,7 +114,9 @@ def forward(self, xin, compound_labels=None):
                     class (float): class index.
                     xc, yc (float) : center of bbox whose values range from 0 to 1.
                     w, h (float) : size of bbox whose values range from 0 to 1.
-        Returns:
+
+        Returns
+        -------
             loss (torch.Tensor): total loss - the target of backprop.
             loss_xy (torch.Tensor): x, y loss - calculated by binary cross entropy (BCE) \
                 with boxsize-dependent weights.
@@ -319,7 +324,6 @@ def __init__(self, config_model, layer_no, in_ch, ignore_thre=0.7):
             in_ch (int): number of input channels.
             ignore_thre (float): threshold of IoU above which objectness training is ignored.
         """
-
         super(YOLOimgLayer, self).__init__()
         strides = [32, 16, 8]  # fixed
         self.anchors = config_model["ANCHORS"]
@@ -356,7 +360,9 @@ def forward(self, xin, all_labels=None):
                     class (float): class index.
                     xc, yc (float) : center of bbox whose values range from 0 to 1.
                     w, h (float) : size of bbox whose values range from 0 to 1.
-        Returns:
+
+        Returns
+        -------
             loss (torch.Tensor): total loss - the target of backprop.
             loss_xy (torch.Tensor): x, y loss - calculated by binary cross entropy (BCE) \
                 with boxsize-dependent weights.

openpmcvl/granular/models/yolov3.py

@@ -1,8 +1,9 @@
+from collections import defaultdict
+
 import torch
-import torch.nn as nn
+from torch import nn
 
-from collections import defaultdict
-from openpmcvl.granular.models.yolo_layer import YOLOLayer, YOLOimgLayer
+from openpmcvl.granular.models.yolo_layer import YOLOimgLayer, YOLOLayer
 
 
 def add_conv(in_ch, out_ch, ksize, stride):
@@ -13,7 +14,9 @@ def add_conv(in_ch, out_ch, ksize, stride):
         out_ch (int): number of output channels of the convolution layer.
         ksize (int): kernel size of the convolution layer.
         stride (int): stride of the convolution layer.
-    Returns:
+
+    Returns
+    -------
         stage (Sequential) : Sequential layers composing a convolution block.
     """
     stage = nn.Sequential()
@@ -71,10 +74,11 @@ def create_yolov3_modules(config_model, ignore_thre):
         config_model (dict): model configuration.
             See YOLOLayer class for details.
         ignore_thre (float): used in YOLOLayer.
-    Returns:
+
+    Returns
+    -------
         mlist (ModuleList): YOLOv3 module list.
     """
-
     # DarkNet53
     mlist = nn.ModuleList()
     mlist.append(add_conv(in_ch=3, out_ch=32, ksize=3, stride=1))
@@ -153,7 +157,8 @@ def forward(self, x, targets=None):
                 where N, C are batchsize and num. of channels.
             targets (torch.Tensor) : label array whose shape is :math:`(N, 50, 5)`
 
-        Returns:
+        Returns
+        -------
             training:
                 output (torch.Tensor): loss tensor for backpropagation.
             test:
@@ -200,10 +205,11 @@ def create_yolov3img_modules(config_model, ignore_thre):
         config_model (dict): model configuration.
             See YOLOLayer class for details.
         ignore_thre (float): used in YOLOLayer.
-    Returns:
+
+    Returns
+    -------
         mlist (ModuleList): YOLOv3 module list.
     """
-
     # DarkNet53
     mlist = nn.ModuleList()
     mlist.append(add_conv(in_ch=4, out_ch=32, ksize=3, stride=1))
@@ -282,7 +288,8 @@ def forward(self, x, targets=None):
                 where N, C are batchsize and num. of channels.
             targets (torch.Tensor) : label array whose shape is :math:`(N, 50, 5)`
 
-        Returns:
+        Returns
+        -------
             training:
                 output (torch.Tensor): loss tensor for backpropagation.
             test:

openpmcvl/granular/pipeline/align.py

@@ -1,8 +1,9 @@
-import os
 import argparse
+import os
 from typing import Dict
 
 from tqdm import tqdm
+
 from openpmcvl.granular.models.subfigure_ocr import classifier
 from openpmcvl.granular.pipeline.utils import load_dataset, save_jsonl
 
@@ -15,7 +16,8 @@ def process_subfigure(model: classifier, subfig_data: Dict) -> Dict:
         model (classifier): Initialized OCR model
         subfig_data (Dict): Dictionary containing subfigure data
 
-    Returns:
+    Returns
+    -------
         Dict: Updated subfigure data with OCR results
     """
     if "subfig_path" not in subfig_data:

openpmcvl/granular/pipeline/classify.py

@@ -1,15 +1,16 @@
 import argparse
-from PIL import Image
 from typing import Any, Dict, List
 
 import torch
-import torch.nn as nn
+from PIL import Image
+from torch import nn
 from torch.utils.data import DataLoader
 from torchvision import models, transforms
 from tqdm import tqdm
 
-from openpmcvl.granular.pipeline.utils import load_dataset, save_jsonl
 from openpmcvl.granular.dataset.dataset import SubfigureDataset
+from openpmcvl.granular.pipeline.utils import load_dataset, save_jsonl
+
 
 MEDICAL_CLASS = 15
 CLASSIFICATION_THRESHOLD = 4
@@ -23,7 +24,8 @@ def load_classification_model(model_path: str, device: torch.device) -> nn.Modul
         model_path (str): Path to the classification model checkpoint
         device (torch.device): Device to use for processing
 
-    Returns:
+    Returns
+    -------
         nn.Module: Loaded classification model
     """
     fig_model = models.resnext101_32x8d()
@@ -55,7 +57,9 @@ def classify_dataset(
         device (torch.device): Device to use for processing.
         output_file (str): Path to save the updated JSONL file with classification results.
         num_workers (int): Number of workers for processing.
-    Returns:
+
+    Returns
+    -------
         None
     """
     mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
@@ -118,7 +122,8 @@ def main(args: argparse.Namespace) -> None:
             - batch_size (int): Batch size for processing
             - output_file (str): Path to save the JSONL file with classification results
 
-    Returns:
+    Returns
+    -------
         None
     """
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")