Add docstrings to functions and reformat (#52)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>

detect.py

@@ -33,6 +33,7 @@
 
 
 def detect(opt):
+    """Detects objects in images using Darknet model, optionally uses a secondary classifier, and performs NMS."""
     if opt.plot_flag:
         os.system("rm -rf " + opt.output_folder + "_img")
         os.makedirs(opt.output_folder + "_img", exist_ok=True)
@@ -228,6 +229,9 @@ def detect(opt):
 
 class ConvNetb(nn.Module):
     def __init__(self, num_classes=60):
+        """Initializes a ConvNetb model with configurable number of classes, defaulting to 60, and a series of
+        convolutional layers.
+        """
         super(ConvNetb, self).__init__()
         n = 64  # initial convolution size
         self.layer1 = nn.Sequential(
@@ -260,6 +264,9 @@ def __init__(self, num_classes=60):
         self.fully_conv = nn.Conv2d(n * 16, 60, kernel_size=4, stride=1, padding=0, bias=True)
 
     def forward(self, x):  # 500 x 1 x 64 x 64
+        """Processes input through 5 layers and a fully connected layer, returning squeezed output; expects input shape
+        500x1x64x64.
+        """
         x = self.layer1(x)
         x = self.layer2(x)
         x = self.layer3(x)

models.py

@@ -72,13 +72,15 @@ class EmptyLayer(nn.Module):
     """Placeholder for 'route' and 'shortcut' layers."""
 
     def __init__(self):
+        """Initializes a placeholder layer for 'route' and 'shortcut' in YOLO architecture."""
         super(EmptyLayer, self).__init__()
 
 
 class YOLOLayer(nn.Module):
     # YOLO Layer 0
 
     def __init__(self, anchors, nC, img_dim, anchor_idxs):
+        """Initializes YOLO layer with given anchors, number of classes, image dimensions, and anchor indexes."""
         super(YOLOLayer, self).__init__()
 
         anchors = [(a_w, a_h) for a_w, a_h in anchors]  # (pixels)
@@ -106,6 +108,7 @@ def __init__(self, anchors, nC, img_dim, anchor_idxs):
         self.anchor_h = self.scaled_anchors[:, 1:2].view((1, nA, 1, 1))
 
     def forward(self, p, targets=None, requestPrecision=False, weight=None, epoch=None):
+        """Processes input tensor `p`, optional targets for precision calculation; returns loss, precision, or both."""
         FT = torch.cuda.FloatTensor if p.is_cuda else torch.FloatTensor
         device = torch.device("cuda:0" if p.is_cuda else "cpu")
         # weight = xview_class_weights(range(60)).to(device)
@@ -226,6 +229,9 @@ class Darknet(nn.Module):
     """YOLOv3 object detection model."""
 
     def __init__(self, config_path, img_size=416):
+        """Initializes Darknet model with a configuration path and optional image size, parsing and creating model
+        modules.
+        """
         super(Darknet, self).__init__()
         self.module_defs = parse_model_config(config_path)
         self.module_defs[0]["height"] = img_size
@@ -234,6 +240,9 @@ def __init__(self, config_path, img_size=416):
         self.loss_names = ["loss", "x", "y", "w", "h", "conf", "cls", "nGT", "TP", "FP", "FPe", "FN", "TC"]
 
     def forward(self, x, targets=None, requestPrecision=False, weight=None, epoch=None):
+        """Processes input through the model, calculates losses, and returns output; includes optional precision
+        computation.
+        """
         is_training = targets is not None
         output = []
         self.losses = defaultdict(float)

scoring/matching.py

@@ -14,9 +14,11 @@
 """
 
 from collections import defaultdict
-from scoring.rectangle import Rectangle
+
 import numpy as np
 
+from scoring.rectangle import Rectangle
+
 
 class Matching(object):
     """Matching class."""
@@ -85,6 +87,7 @@ def _compute_iou_from_rectangle_pairs(self):
             self.iou_matrix = np.zeros((n, m))
 
     def greedy_match(self, iou_threshold):
+        """Performs greedy matching of rectangles based on IOU threshold, returning matched indices."""
         gt_rects_matched = [False for gt_index in range(self.m)]
         rects_matched = [False for r_index in range(self.n)]
 

scoring/score.py

@@ -17,10 +17,9 @@
 import json
 import os
 import time
-import scipy.io
-
 
 import numpy as np
+import scipy.io
 from tqdm import tqdm
 
 from scoring.matching import Matching

train.py

@@ -24,6 +24,7 @@
 
 
 def main(opt):
+    """Initializes and trains a Darknet model for object detection with configurable parameters and data paths."""
     os.makedirs("weights", exist_ok=True)
     cuda = torch.cuda.is_available()
     device = torch.device("cuda:0" if cuda else "cpu")

utils/datasets.py

@@ -9,11 +9,14 @@
 import torch
 
 # from torch.utils.data import Dataset
-from utils.utils import xyxy2xywh, xview_class_weights
+from utils.utils import xview_class_weights, xyxy2xywh
 
 
 class ImageFolder:  # for eval-only
     def __init__(self, path, batch_size=1, img_size=416):
+        """Initializes an ImageFolder object to load images from directory or file for evaluation, with customizable
+        batch size and image size.
+        """
         if os.path.isdir(path):
             self.files = sorted(glob.glob("%s/*.*" % path))
         elif os.path.isfile(path):
@@ -30,10 +33,12 @@ def __init__(self, path, batch_size=1, img_size=416):
         self.rgb_std = np.array([29.99, 24.498, 22.046], dtype=np.float32).reshape((3, 1, 1))
 
     def __iter__(self):
+        """Initializes iterator by resetting count and returns the iterator object itself for sequential access."""
         self.count = -1
         return self
 
     def __next__(self):
+        """Iterates to the next image path in the dataset, raising StopIteration when all images are iterated."""
         self.count += 1
         if self.count == self.nB:
             raise StopIteration
@@ -51,11 +56,15 @@ def __next__(self):
         return [img_path], img
 
     def __len__(self):
+        """Returns the number of batches in the dataset."""
         return self.nB  # number of batches
 
 
 class ListDataset:  # for training
     def __init__(self, path, batch_size=1, img_size=608, targets_path=""):
+        """Initializes ListDataset for image training with optional batch size and target path, ensuring image path
+        contains images.
+        """
         self.path = path
         self.files = sorted(glob.glob("%s/*.tif" % path))
         self.nF = len(self.files)  # number of image files
@@ -88,6 +97,9 @@ def __init__(self, path, batch_size=1, img_size=608, targets_path=""):
         # self.rgb_std = np.array([69.095, 66.369, 64.236], dtype=np.float32).reshape((1, 3, 1, 1))
 
     def __iter__(self):
+        """Initializes iterator by resetting count, creating a shuffled vector of image numbers based on their
+        weights.
+        """
         self.count = -1
         # self.shuffled_vector = np.random.permutation(self.nF)  # shuffled vector
         self.shuffled_vector = np.random.choice(
@@ -97,6 +109,7 @@ def __iter__(self):
 
     # @profile
     def __next__(self):
+        """Advances to the next batch of data, raising StopIteration when the dataset end is reached."""
         self.count += 1
         if self.count == self.nB:
             raise StopIteration
@@ -268,10 +281,12 @@ def __next__(self):
         return torch.from_numpy(img_all), labels_all
 
     def __len__(self):
+        """Returns the number of batches in the dataset."""
         return self.nB  # number of batches
 
 
 def resize_square(img, height=416, color=(0, 0, 0)):  # resizes a rectangular image to a padded square
+    """Resizes an image to a padded square of given height, maintaining aspect ratio; default color is black."""
     shape = img.shape[:2]  # shape = [height, width]
     ratio = float(height) / max(shape)
     new_shape = [round(shape[0] * ratio), round(shape[1] * ratio)]
@@ -358,7 +373,9 @@ def random_affine(
 
 
 def convert_tif2bmp(p="/Users/glennjocher/Downloads/DATA/xview/val_images_bmp"):
+    """Converts TIF images to BMP format in a specified path, deleting the original TIF files."""
     import glob
+
     import cv2
 
     files = sorted(glob.glob("%s/*.tif" % p))

utils/utils.py

@@ -18,6 +18,7 @@ def load_classes(path):
 
 
 def modelinfo(model):
+    """Prints model layers, parameters, gradients, and statistics; requires a model object as input."""
     nparams = sum(x.numel() for x in model.parameters())
     ngradients = sum(x.numel() for x in model.parameters() if x.requires_grad)
     print("\n%4s %70s %9s %12s %20s %12s %12s" % ("", "name", "gradient", "parameters", "shape", "mu", "sigma"))
@@ -30,6 +31,7 @@ def modelinfo(model):
 
 
 def xview_classes2indices(classes):  # remap xview classes 11-94 to 0-61
+    """Remaps xview classes (11-94) to indices (0-61), skipping unassigned classes; classes not mapped are set to -1."""
     indices = [
         -1,
         -1,
@@ -131,6 +133,7 @@ def xview_classes2indices(classes):  # remap xview classes 11-94 to 0-61
 
 
 def xview_indices2classes(indices):  # remap xview classes 11-94 to 0-61
+    """Remaps xView dataset class indices (11-94) to a contiguous range (0-61)."""
     class_list = [
         11,
         12,
@@ -197,6 +200,7 @@ def xview_indices2classes(indices):  # remap xview classes 11-94 to 0-61
 
 
 def xview_class_weights(indices):  # weights of each class in the training set, normalized to mu = 1
+    """Calculates normalized class weights from given indices, with mean=1, for class imbalance handling."""
     weights = 1 / torch.FloatTensor(
         [
             74,
@@ -266,6 +270,7 @@ def xview_class_weights(indices):  # weights of each class in the training set,
 
 
 def xview_class_weights_hard_mining(indices):  # weights of each class in the training set, normalized to mu = 1
+    """Calculates normalized class weights for hard-mined classes, useful for imbalanced datasets."""
     weights = 1 / torch.FloatTensor(
         [
             33.97268,
@@ -335,6 +340,7 @@ def xview_class_weights_hard_mining(indices):  # weights of each class in the tr
 
 
 def plot_one_box(x, im, color=None, label=None, line_thickness=None):
+    """Draws a labeled rectangle with specified thickness and color on an image."""
     tl = line_thickness or round(0.003 * max(im.shape[0:2]))  # line thickness
     color = color or [random.randint(0, 255) for _ in range(3)]
     c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
@@ -348,6 +354,7 @@ def plot_one_box(x, im, color=None, label=None, line_thickness=None):
 
 
 def weights_init_normal(m):
+    """Initializes network weights normally for Conv and BatchNorm2d layers."""
     classname = m.__class__.__name__
     if classname.find("Conv") != -1:
         torch.nn.init.normal_(m.weight.data, 0.0, 0.03)
@@ -357,6 +364,7 @@ def weights_init_normal(m):
 
 
 def xyxy2xywh(box):
+    """Converts bounding box format from [x1, y1, x2, y2] to [x_center, y_center, width, height]."""
     xywh = np.zeros(box.shape)
     xywh[:, 0] = (box[:, 0] + box[:, 2]) / 2
     xywh[:, 1] = (box[:, 1] + box[:, 3]) / 2
@@ -658,7 +666,9 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4, mat=None, img
 
 
 def secondary_class_detection(x, y, w, h, img, model, device):
-    # Runs secondary classifier on bounding boxes
+    """Applies secondary classification to bounding boxes using a given model and returns the class with the highest
+    probability for each box.
+    """
     print("Classifying boxes...", end="")
 
     # 1. create 48-pixel squares from each chip
@@ -713,13 +723,14 @@ def secondary_class_detection(x, y, w, h, img, model, device):
 
 
 def createChips():
-    # Creates *.h5 file of all chips in xview dataset for training independent classifier
+    """Generates and saves a dataset of image chips from the xview dataset for classifier training."""
+
+    from sys import platform
 
-    import scipy.io
-    import numpy as np
     import cv2
     import h5py
-    from sys import platform
+    import numpy as np
+    import scipy.io
 
     mat = scipy.io.loadmat("utils/targets_c60.mat")
     unique_images = np.unique(mat["id"])
@@ -769,7 +780,9 @@ def createChips():
 
 
 def strip_optimizer_from_checkpoint(filename="weights/best.pt"):
-    # Strip optimizer from *.pt files for lighter files (reduced by 2/3 size)
+    """Strips optimizer from .pt checkpoint files, reducing size by 2/3, by saving a lite version without optimizer
+    state.
+    """
     import torch
 
     a = torch.load(filename, map_location="cpu")
@@ -778,9 +791,9 @@ def strip_optimizer_from_checkpoint(filename="weights/best.pt"):
 
 
 def plotResults():
-    # Plot YOLO training results
-    import numpy as np
+    """Plots YOLO training results from 'results.txt' for key metrics over first 300 epochs."""
     import matplotlib.pyplot as plt
+    import numpy as np
 
     plt.figure(figsize=(16, 8))
     s = ["X", "Y", "Width", "Height", "Objectness", "Classification", "Total Loss", "Precision", "Recall"]

utils/utils_xview.py

@@ -5,13 +5,14 @@
 
 
 def xview_class2name(classes):
+    """Converts numerical class IDs to their corresponding names using 'data/xview.names'."""
     with open("data/xview.names", "r") as f:
         x = f.readlines()
     return x[classes].replace("\n", "")
 
 
 def get_labels(fname="xView_train.geojson"):  # from utils.utils_xview import get_labels; get_labels()
-    # Official function supplied by xView
+    """Parses 'xView_train.geojson' to extract object coordinates, chip IDs, and class IDs as numpy arrays."""
     with open(fname) as f:
         data = json.load(f)
 
@@ -36,10 +37,10 @@ def get_labels(fname="xView_train.geojson"):  # from utils.utils_xview import ge
 
 
 def create_mat_file():
-    # saves geojson file to .mat format for analysis in MATLAB
-    import scipy.io
+    """Saves geojson data as a MATLAB (.mat) file, enriching it with image statistics and shapes."""
     import cv2
     import numpy as np
+    import scipy.io
 
     path = "/Users/glennjocher/Downloads/DATA/xview/"
     coords, chips, classes = get_labels(path + "xView_train.geojson")