Add docstrings and improve comments (#11229)
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
parent
ccfc1cf925
commit
d5458f27cd
16 changed files with 34 additions and 17 deletions
|
|
@ -67,6 +67,7 @@ server.login(from_email, password)
|
|||
|
||||
```python
|
||||
def send_email(to_email, from_email, object_detected=1):
|
||||
"""Sends an email notification indicating the number of objects detected; defaults to 1 object."""
|
||||
message = MIMEMultipart()
|
||||
message['From'] = from_email
|
||||
message['To'] = to_email
|
||||
|
|
@ -83,7 +84,7 @@ def send_email(to_email, from_email, object_detected=1):
|
|||
```python
|
||||
class ObjectDetection:
|
||||
def __init__(self, capture_index):
|
||||
# default parameters
|
||||
"""Initializes an ObjectDetection instance with a given camera index."""
|
||||
self.capture_index = capture_index
|
||||
self.email_sent = False
|
||||
|
||||
|
|
@ -99,10 +100,12 @@ class ObjectDetection:
|
|||
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
|
||||
def predict(self, im0):
|
||||
"""Run prediction using a YOLO model for the input image `im0`."""
|
||||
results = self.model(im0)
|
||||
return results
|
||||
|
||||
def display_fps(self, im0):
|
||||
"""Displays the FPS on an image `im0` by calculating and overlaying as white text on a black rectangle."""
|
||||
self.end_time = time()
|
||||
fps = 1 / np.round(self.end_time - self.start_time, 2)
|
||||
text = f'FPS: {int(fps)}'
|
||||
|
|
@ -112,6 +115,7 @@ class ObjectDetection:
|
|||
cv2.putText(im0, text, (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2)
|
||||
|
||||
def plot_bboxes(self, results, im0):
|
||||
"""Plots bounding boxes on an image given detection results; returns annotated image and class IDs."""
|
||||
class_ids = []
|
||||
self.annotator = Annotator(im0, 3, results[0].names)
|
||||
boxes = results[0].boxes.xyxy.cpu()
|
||||
|
|
@ -123,6 +127,7 @@ class ObjectDetection:
|
|||
return im0, class_ids
|
||||
|
||||
def __call__(self):
|
||||
"""Executes object detection on video frames from a specified camera index, plotting bounding boxes and returning modified frames."""
|
||||
cap = cv2.VideoCapture(self.capture_index)
|
||||
assert cap.isOpened()
|
||||
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ shared_model = YOLO("yolov8n.pt")
|
|||
|
||||
|
||||
def predict(image_path):
|
||||
"""Predicts objects in an image using a preloaded YOLO model, take path string to image as argument."""
|
||||
results = shared_model.predict(image_path)
|
||||
# Process results
|
||||
|
||||
|
|
@ -62,6 +63,7 @@ shared_model_2 = YOLO("yolov8n_2.pt")
|
|||
|
||||
|
||||
def predict(model, image_path):
|
||||
"""Runs prediction on an image using a specified YOLO model, returning the results."""
|
||||
results = model.predict(image_path)
|
||||
# Process results
|
||||
|
||||
|
|
@ -88,7 +90,7 @@ from threading import Thread
|
|||
|
||||
|
||||
def thread_safe_predict(image_path):
|
||||
# Instantiate a new model inside the thread
|
||||
"""Predict on an image using a new YOLO model instance in a thread-safe manner; takes image path as input."""
|
||||
local_model = YOLO("yolov8n.pt")
|
||||
results = local_model.predict(image_path)
|
||||
# Process results
|
||||
|
|
|
|||
|
|
@ -118,6 +118,7 @@ After creating the AWS CloudFormation Stack, the next step is to deploy YOLOv8.
|
|||
import json
|
||||
|
||||
def output_fn(prediction_output, content_type):
|
||||
"""Formats model outputs as JSON string according to content_type, extracting attributes like boxes, masks, keypoints."""
|
||||
print("Executing output_fn from inference.py ...")
|
||||
infer = {}
|
||||
for result in prediction_output:
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ model = YOLO("yolov8n.pt")
|
|||
|
||||
|
||||
def predict_image(img, conf_threshold, iou_threshold):
|
||||
"""Predicts and plots labeled objects in an image using YOLOv8 model with adjustable confidence and IOU thresholds."""
|
||||
results = model.predict(
|
||||
source=img,
|
||||
conf=conf_threshold,
|
||||
|
|
|
|||
|
|
@ -731,7 +731,7 @@ When using YOLO models in a multi-threaded application, it's important to instan
|
|||
from threading import Thread
|
||||
|
||||
def thread_safe_predict(image_path):
|
||||
# Instantiate a new model inside the thread
|
||||
"""Performs thread-safe prediction on an image using a locally instantiated YOLO model."""
|
||||
local_model = YOLO("yolov8n.pt")
|
||||
results = local_model.predict(image_path)
|
||||
# Process results
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ from ultralytics import YOLO
|
|||
|
||||
|
||||
def on_predict_batch_end(predictor):
|
||||
# Retrieve the batch data
|
||||
"""Handle prediction batch end by combining results with corresponding frames; modifies predictor results."""
|
||||
_, image, _, _ = predictor.batch
|
||||
|
||||
# Ensure that image is a list
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ from ultralytics.models.yolo.detect import DetectionTrainer
|
|||
|
||||
class CustomTrainer(DetectionTrainer):
|
||||
def get_model(self, cfg, weights):
|
||||
"""Loads a custom detection model given configuration and weight files."""
|
||||
...
|
||||
|
||||
|
||||
|
|
@ -65,16 +66,19 @@ from ultralytics.nn.tasks import DetectionModel
|
|||
|
||||
class MyCustomModel(DetectionModel):
|
||||
def init_criterion(self):
|
||||
"""Initializes the loss function and adds a callback for uploading the model to Google Drive every 10 epochs."""
|
||||
...
|
||||
|
||||
|
||||
class CustomTrainer(DetectionTrainer):
|
||||
def get_model(self, cfg, weights):
|
||||
"""Returns a customized detection model instance configured with specified config and weights."""
|
||||
return MyCustomModel(...)
|
||||
|
||||
|
||||
# callback to upload model weights
|
||||
def log_model(trainer):
|
||||
"""Logs the path of the last model weight used by the trainer."""
|
||||
last_weight_path = trainer.last
|
||||
print(last_weight_path)
|
||||
|
||||
|
|
|
|||
|
|
@ -38,12 +38,14 @@ import torch.nn as nn
|
|||
|
||||
class SPP(nn.Module):
|
||||
def __init__(self):
|
||||
"""Initializes an SPP module with three different sizes of max pooling layers."""
|
||||
super().__init__()
|
||||
self.maxpool1 = nn.MaxPool2d(5, 1, padding=2)
|
||||
self.maxpool2 = nn.MaxPool2d(9, 1, padding=4)
|
||||
self.maxpool3 = nn.MaxPool2d(13, 1, padding=6)
|
||||
|
||||
def forward(self, x):
|
||||
"""Applies three max pooling layers on input `x` and concatenates results along channel dimension."""
|
||||
o1 = self.maxpool1(x)
|
||||
o2 = self.maxpool2(x)
|
||||
o3 = self.maxpool3(x)
|
||||
|
|
@ -52,10 +54,12 @@ class SPP(nn.Module):
|
|||
|
||||
class SPPF(nn.Module):
|
||||
def __init__(self):
|
||||
"""Initializes an SPPF module with a specific configuration of MaxPool2d layer."""
|
||||
super().__init__()
|
||||
self.maxpool = nn.MaxPool2d(5, 1, padding=2)
|
||||
|
||||
def forward(self, x):
|
||||
"""Applies sequential max pooling and concatenates results with input tensor; expects input tensor x of any shape."""
|
||||
o1 = self.maxpool(x)
|
||||
o2 = self.maxpool(o1)
|
||||
o3 = self.maxpool(o2)
|
||||
|
|
@ -63,6 +67,7 @@ class SPPF(nn.Module):
|
|||
|
||||
|
||||
def main():
|
||||
"""Compares outputs and performance of SPP and SPPF on a random tensor (8, 32, 16, 16)."""
|
||||
input_tensor = torch.rand(8, 32, 16, 16)
|
||||
spp = SPP()
|
||||
sppf = SPPF()
|
||||
|
|
|
|||
|
|
@ -65,7 +65,7 @@ Fitness is the value we seek to maximize. In YOLOv5 we define a default fitness
|
|||
|
||||
```python
|
||||
def fitness(x):
|
||||
# Model fitness as a weighted combination of metrics
|
||||
"""Evaluates model fitness by summing weighted metrics [P, R, mAP@0.5, mAP@0.5:0.95], x is a numpy array of shape (n, 4)."""
|
||||
w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
|
||||
return (x[:, :4] * w).sum(1)
|
||||
```
|
||||
|
|
|
|||
|
|
@ -179,6 +179,7 @@ import threading
|
|||
|
||||
|
||||
def run(model, im):
|
||||
"""Performs inference on an image using a given model and saves the output; model must support `.save()` method."""
|
||||
results = model(im)
|
||||
results.save()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue