Add docstrings and improve comments (#11229)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-05-03 22:51:26 +02:00 committed by GitHub
parent ccfc1cf925
commit d5458f27cd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 34 additions and 17 deletions

View file

@ -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)

View file

@ -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