ultralytics-opencv-headless 8.3.242__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tests/__init__.py +23 -0
- tests/conftest.py +59 -0
- tests/test_cli.py +131 -0
- tests/test_cuda.py +216 -0
- tests/test_engine.py +157 -0
- tests/test_exports.py +309 -0
- tests/test_integrations.py +151 -0
- tests/test_python.py +777 -0
- tests/test_solutions.py +371 -0
- ultralytics/__init__.py +48 -0
- ultralytics/assets/bus.jpg +0 -0
- ultralytics/assets/zidane.jpg +0 -0
- ultralytics/cfg/__init__.py +1026 -0
- ultralytics/cfg/datasets/Argoverse.yaml +78 -0
- ultralytics/cfg/datasets/DOTAv1.5.yaml +37 -0
- ultralytics/cfg/datasets/DOTAv1.yaml +36 -0
- ultralytics/cfg/datasets/GlobalWheat2020.yaml +68 -0
- ultralytics/cfg/datasets/HomeObjects-3K.yaml +32 -0
- ultralytics/cfg/datasets/ImageNet.yaml +2025 -0
- ultralytics/cfg/datasets/Objects365.yaml +447 -0
- ultralytics/cfg/datasets/SKU-110K.yaml +58 -0
- ultralytics/cfg/datasets/VOC.yaml +102 -0
- ultralytics/cfg/datasets/VisDrone.yaml +87 -0
- ultralytics/cfg/datasets/african-wildlife.yaml +25 -0
- ultralytics/cfg/datasets/brain-tumor.yaml +22 -0
- ultralytics/cfg/datasets/carparts-seg.yaml +44 -0
- ultralytics/cfg/datasets/coco-pose.yaml +64 -0
- ultralytics/cfg/datasets/coco.yaml +118 -0
- ultralytics/cfg/datasets/coco128-seg.yaml +101 -0
- ultralytics/cfg/datasets/coco128.yaml +101 -0
- ultralytics/cfg/datasets/coco8-grayscale.yaml +103 -0
- ultralytics/cfg/datasets/coco8-multispectral.yaml +104 -0
- ultralytics/cfg/datasets/coco8-pose.yaml +47 -0
- ultralytics/cfg/datasets/coco8-seg.yaml +101 -0
- ultralytics/cfg/datasets/coco8.yaml +101 -0
- ultralytics/cfg/datasets/construction-ppe.yaml +32 -0
- ultralytics/cfg/datasets/crack-seg.yaml +22 -0
- ultralytics/cfg/datasets/dog-pose.yaml +52 -0
- ultralytics/cfg/datasets/dota8-multispectral.yaml +38 -0
- ultralytics/cfg/datasets/dota8.yaml +35 -0
- ultralytics/cfg/datasets/hand-keypoints.yaml +50 -0
- ultralytics/cfg/datasets/kitti.yaml +27 -0
- ultralytics/cfg/datasets/lvis.yaml +1240 -0
- ultralytics/cfg/datasets/medical-pills.yaml +21 -0
- ultralytics/cfg/datasets/open-images-v7.yaml +663 -0
- ultralytics/cfg/datasets/package-seg.yaml +22 -0
- ultralytics/cfg/datasets/signature.yaml +21 -0
- ultralytics/cfg/datasets/tiger-pose.yaml +41 -0
- ultralytics/cfg/datasets/xView.yaml +155 -0
- ultralytics/cfg/default.yaml +130 -0
- ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +17 -0
- ultralytics/cfg/models/11/yolo11-cls.yaml +33 -0
- ultralytics/cfg/models/11/yolo11-obb.yaml +50 -0
- ultralytics/cfg/models/11/yolo11-pose.yaml +51 -0
- ultralytics/cfg/models/11/yolo11-seg.yaml +50 -0
- ultralytics/cfg/models/11/yolo11.yaml +50 -0
- ultralytics/cfg/models/11/yoloe-11-seg.yaml +48 -0
- ultralytics/cfg/models/11/yoloe-11.yaml +48 -0
- ultralytics/cfg/models/12/yolo12-cls.yaml +32 -0
- ultralytics/cfg/models/12/yolo12-obb.yaml +48 -0
- ultralytics/cfg/models/12/yolo12-pose.yaml +49 -0
- ultralytics/cfg/models/12/yolo12-seg.yaml +48 -0
- ultralytics/cfg/models/12/yolo12.yaml +48 -0
- ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +53 -0
- ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +45 -0
- ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +45 -0
- ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +57 -0
- ultralytics/cfg/models/v10/yolov10b.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10l.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10m.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10n.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10s.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10x.yaml +45 -0
- ultralytics/cfg/models/v3/yolov3-spp.yaml +49 -0
- ultralytics/cfg/models/v3/yolov3-tiny.yaml +40 -0
- ultralytics/cfg/models/v3/yolov3.yaml +49 -0
- ultralytics/cfg/models/v5/yolov5-p6.yaml +62 -0
- ultralytics/cfg/models/v5/yolov5.yaml +51 -0
- ultralytics/cfg/models/v6/yolov6.yaml +56 -0
- ultralytics/cfg/models/v8/yoloe-v8-seg.yaml +48 -0
- ultralytics/cfg/models/v8/yoloe-v8.yaml +48 -0
- ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +28 -0
- ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +28 -0
- ultralytics/cfg/models/v8/yolov8-cls.yaml +32 -0
- ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +58 -0
- ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +60 -0
- ultralytics/cfg/models/v8/yolov8-ghost.yaml +50 -0
- ultralytics/cfg/models/v8/yolov8-obb.yaml +49 -0
- ultralytics/cfg/models/v8/yolov8-p2.yaml +57 -0
- ultralytics/cfg/models/v8/yolov8-p6.yaml +59 -0
- ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +60 -0
- ultralytics/cfg/models/v8/yolov8-pose.yaml +50 -0
- ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +49 -0
- ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +59 -0
- ultralytics/cfg/models/v8/yolov8-seg.yaml +49 -0
- ultralytics/cfg/models/v8/yolov8-world.yaml +51 -0
- ultralytics/cfg/models/v8/yolov8-worldv2.yaml +49 -0
- ultralytics/cfg/models/v8/yolov8.yaml +49 -0
- ultralytics/cfg/models/v9/yolov9c-seg.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9c.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9e-seg.yaml +64 -0
- ultralytics/cfg/models/v9/yolov9e.yaml +64 -0
- ultralytics/cfg/models/v9/yolov9m.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9s.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9t.yaml +41 -0
- ultralytics/cfg/trackers/botsort.yaml +21 -0
- ultralytics/cfg/trackers/bytetrack.yaml +12 -0
- ultralytics/data/__init__.py +26 -0
- ultralytics/data/annotator.py +66 -0
- ultralytics/data/augment.py +2801 -0
- ultralytics/data/base.py +435 -0
- ultralytics/data/build.py +437 -0
- ultralytics/data/converter.py +855 -0
- ultralytics/data/dataset.py +834 -0
- ultralytics/data/loaders.py +704 -0
- ultralytics/data/scripts/download_weights.sh +18 -0
- ultralytics/data/scripts/get_coco.sh +61 -0
- ultralytics/data/scripts/get_coco128.sh +18 -0
- ultralytics/data/scripts/get_imagenet.sh +52 -0
- ultralytics/data/split.py +138 -0
- ultralytics/data/split_dota.py +344 -0
- ultralytics/data/utils.py +798 -0
- ultralytics/engine/__init__.py +1 -0
- ultralytics/engine/exporter.py +1574 -0
- ultralytics/engine/model.py +1124 -0
- ultralytics/engine/predictor.py +508 -0
- ultralytics/engine/results.py +1522 -0
- ultralytics/engine/trainer.py +974 -0
- ultralytics/engine/tuner.py +448 -0
- ultralytics/engine/validator.py +384 -0
- ultralytics/hub/__init__.py +166 -0
- ultralytics/hub/auth.py +151 -0
- ultralytics/hub/google/__init__.py +174 -0
- ultralytics/hub/session.py +422 -0
- ultralytics/hub/utils.py +162 -0
- ultralytics/models/__init__.py +9 -0
- ultralytics/models/fastsam/__init__.py +7 -0
- ultralytics/models/fastsam/model.py +79 -0
- ultralytics/models/fastsam/predict.py +169 -0
- ultralytics/models/fastsam/utils.py +23 -0
- ultralytics/models/fastsam/val.py +38 -0
- ultralytics/models/nas/__init__.py +7 -0
- ultralytics/models/nas/model.py +98 -0
- ultralytics/models/nas/predict.py +56 -0
- ultralytics/models/nas/val.py +38 -0
- ultralytics/models/rtdetr/__init__.py +7 -0
- ultralytics/models/rtdetr/model.py +63 -0
- ultralytics/models/rtdetr/predict.py +88 -0
- ultralytics/models/rtdetr/train.py +89 -0
- ultralytics/models/rtdetr/val.py +216 -0
- ultralytics/models/sam/__init__.py +25 -0
- ultralytics/models/sam/amg.py +275 -0
- ultralytics/models/sam/build.py +365 -0
- ultralytics/models/sam/build_sam3.py +377 -0
- ultralytics/models/sam/model.py +169 -0
- ultralytics/models/sam/modules/__init__.py +1 -0
- ultralytics/models/sam/modules/blocks.py +1067 -0
- ultralytics/models/sam/modules/decoders.py +495 -0
- ultralytics/models/sam/modules/encoders.py +794 -0
- ultralytics/models/sam/modules/memory_attention.py +298 -0
- ultralytics/models/sam/modules/sam.py +1160 -0
- ultralytics/models/sam/modules/tiny_encoder.py +979 -0
- ultralytics/models/sam/modules/transformer.py +344 -0
- ultralytics/models/sam/modules/utils.py +512 -0
- ultralytics/models/sam/predict.py +3940 -0
- ultralytics/models/sam/sam3/__init__.py +3 -0
- ultralytics/models/sam/sam3/decoder.py +546 -0
- ultralytics/models/sam/sam3/encoder.py +529 -0
- ultralytics/models/sam/sam3/geometry_encoders.py +415 -0
- ultralytics/models/sam/sam3/maskformer_segmentation.py +286 -0
- ultralytics/models/sam/sam3/model_misc.py +199 -0
- ultralytics/models/sam/sam3/necks.py +129 -0
- ultralytics/models/sam/sam3/sam3_image.py +339 -0
- ultralytics/models/sam/sam3/text_encoder_ve.py +307 -0
- ultralytics/models/sam/sam3/vitdet.py +547 -0
- ultralytics/models/sam/sam3/vl_combiner.py +160 -0
- ultralytics/models/utils/__init__.py +1 -0
- ultralytics/models/utils/loss.py +466 -0
- ultralytics/models/utils/ops.py +315 -0
- ultralytics/models/yolo/__init__.py +7 -0
- ultralytics/models/yolo/classify/__init__.py +7 -0
- ultralytics/models/yolo/classify/predict.py +90 -0
- ultralytics/models/yolo/classify/train.py +202 -0
- ultralytics/models/yolo/classify/val.py +216 -0
- ultralytics/models/yolo/detect/__init__.py +7 -0
- ultralytics/models/yolo/detect/predict.py +122 -0
- ultralytics/models/yolo/detect/train.py +227 -0
- ultralytics/models/yolo/detect/val.py +507 -0
- ultralytics/models/yolo/model.py +430 -0
- ultralytics/models/yolo/obb/__init__.py +7 -0
- ultralytics/models/yolo/obb/predict.py +56 -0
- ultralytics/models/yolo/obb/train.py +79 -0
- ultralytics/models/yolo/obb/val.py +302 -0
- ultralytics/models/yolo/pose/__init__.py +7 -0
- ultralytics/models/yolo/pose/predict.py +65 -0
- ultralytics/models/yolo/pose/train.py +110 -0
- ultralytics/models/yolo/pose/val.py +248 -0
- ultralytics/models/yolo/segment/__init__.py +7 -0
- ultralytics/models/yolo/segment/predict.py +109 -0
- ultralytics/models/yolo/segment/train.py +69 -0
- ultralytics/models/yolo/segment/val.py +307 -0
- ultralytics/models/yolo/world/__init__.py +5 -0
- ultralytics/models/yolo/world/train.py +173 -0
- ultralytics/models/yolo/world/train_world.py +178 -0
- ultralytics/models/yolo/yoloe/__init__.py +22 -0
- ultralytics/models/yolo/yoloe/predict.py +162 -0
- ultralytics/models/yolo/yoloe/train.py +287 -0
- ultralytics/models/yolo/yoloe/train_seg.py +122 -0
- ultralytics/models/yolo/yoloe/val.py +206 -0
- ultralytics/nn/__init__.py +27 -0
- ultralytics/nn/autobackend.py +958 -0
- ultralytics/nn/modules/__init__.py +182 -0
- ultralytics/nn/modules/activation.py +54 -0
- ultralytics/nn/modules/block.py +1947 -0
- ultralytics/nn/modules/conv.py +669 -0
- ultralytics/nn/modules/head.py +1183 -0
- ultralytics/nn/modules/transformer.py +793 -0
- ultralytics/nn/modules/utils.py +159 -0
- ultralytics/nn/tasks.py +1768 -0
- ultralytics/nn/text_model.py +356 -0
- ultralytics/py.typed +1 -0
- ultralytics/solutions/__init__.py +41 -0
- ultralytics/solutions/ai_gym.py +108 -0
- ultralytics/solutions/analytics.py +264 -0
- ultralytics/solutions/config.py +107 -0
- ultralytics/solutions/distance_calculation.py +123 -0
- ultralytics/solutions/heatmap.py +125 -0
- ultralytics/solutions/instance_segmentation.py +86 -0
- ultralytics/solutions/object_blurrer.py +89 -0
- ultralytics/solutions/object_counter.py +190 -0
- ultralytics/solutions/object_cropper.py +87 -0
- ultralytics/solutions/parking_management.py +280 -0
- ultralytics/solutions/queue_management.py +93 -0
- ultralytics/solutions/region_counter.py +133 -0
- ultralytics/solutions/security_alarm.py +151 -0
- ultralytics/solutions/similarity_search.py +219 -0
- ultralytics/solutions/solutions.py +828 -0
- ultralytics/solutions/speed_estimation.py +114 -0
- ultralytics/solutions/streamlit_inference.py +260 -0
- ultralytics/solutions/templates/similarity-search.html +156 -0
- ultralytics/solutions/trackzone.py +88 -0
- ultralytics/solutions/vision_eye.py +67 -0
- ultralytics/trackers/__init__.py +7 -0
- ultralytics/trackers/basetrack.py +115 -0
- ultralytics/trackers/bot_sort.py +257 -0
- ultralytics/trackers/byte_tracker.py +469 -0
- ultralytics/trackers/track.py +116 -0
- ultralytics/trackers/utils/__init__.py +1 -0
- ultralytics/trackers/utils/gmc.py +339 -0
- ultralytics/trackers/utils/kalman_filter.py +482 -0
- ultralytics/trackers/utils/matching.py +154 -0
- ultralytics/utils/__init__.py +1450 -0
- ultralytics/utils/autobatch.py +118 -0
- ultralytics/utils/autodevice.py +205 -0
- ultralytics/utils/benchmarks.py +728 -0
- ultralytics/utils/callbacks/__init__.py +5 -0
- ultralytics/utils/callbacks/base.py +233 -0
- ultralytics/utils/callbacks/clearml.py +146 -0
- ultralytics/utils/callbacks/comet.py +625 -0
- ultralytics/utils/callbacks/dvc.py +197 -0
- ultralytics/utils/callbacks/hub.py +110 -0
- ultralytics/utils/callbacks/mlflow.py +134 -0
- ultralytics/utils/callbacks/neptune.py +126 -0
- ultralytics/utils/callbacks/platform.py +73 -0
- ultralytics/utils/callbacks/raytune.py +42 -0
- ultralytics/utils/callbacks/tensorboard.py +123 -0
- ultralytics/utils/callbacks/wb.py +188 -0
- ultralytics/utils/checks.py +998 -0
- ultralytics/utils/cpu.py +85 -0
- ultralytics/utils/dist.py +123 -0
- ultralytics/utils/downloads.py +529 -0
- ultralytics/utils/errors.py +35 -0
- ultralytics/utils/events.py +113 -0
- ultralytics/utils/export/__init__.py +7 -0
- ultralytics/utils/export/engine.py +237 -0
- ultralytics/utils/export/imx.py +315 -0
- ultralytics/utils/export/tensorflow.py +231 -0
- ultralytics/utils/files.py +219 -0
- ultralytics/utils/git.py +137 -0
- ultralytics/utils/instance.py +484 -0
- ultralytics/utils/logger.py +444 -0
- ultralytics/utils/loss.py +849 -0
- ultralytics/utils/metrics.py +1560 -0
- ultralytics/utils/nms.py +337 -0
- ultralytics/utils/ops.py +664 -0
- ultralytics/utils/patches.py +201 -0
- ultralytics/utils/plotting.py +1045 -0
- ultralytics/utils/tal.py +403 -0
- ultralytics/utils/torch_utils.py +984 -0
- ultralytics/utils/tqdm.py +440 -0
- ultralytics/utils/triton.py +112 -0
- ultralytics/utils/tuner.py +160 -0
- ultralytics_opencv_headless-8.3.242.dist-info/METADATA +374 -0
- ultralytics_opencv_headless-8.3.242.dist-info/RECORD +298 -0
- ultralytics_opencv_headless-8.3.242.dist-info/WHEEL +5 -0
- ultralytics_opencv_headless-8.3.242.dist-info/entry_points.txt +3 -0
- ultralytics_opencv_headless-8.3.242.dist-info/licenses/LICENSE +661 -0
- ultralytics_opencv_headless-8.3.242.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,114 @@
|
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
2
|
+
|
|
3
|
+
from collections import deque
|
|
4
|
+
from math import sqrt
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
|
|
8
|
+
from ultralytics.utils.plotting import colors
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class SpeedEstimator(BaseSolution):
|
|
12
|
+
"""A class to estimate the speed of objects in a real-time video stream based on their tracks.
|
|
13
|
+
|
|
14
|
+
This class extends the BaseSolution class and provides functionality for estimating object speeds using tracking
|
|
15
|
+
data in video streams. Speed is calculated based on pixel displacement over time and converted to real-world units
|
|
16
|
+
using a configurable meters-per-pixel scale factor.
|
|
17
|
+
|
|
18
|
+
Attributes:
|
|
19
|
+
fps (float): Video frame rate for time calculations.
|
|
20
|
+
frame_count (int): Global frame counter for tracking temporal information.
|
|
21
|
+
trk_frame_ids (dict): Maps track IDs to their first frame index.
|
|
22
|
+
spd (dict): Final speed per object in km/h once locked.
|
|
23
|
+
trk_hist (dict): Maps track IDs to deque of position history.
|
|
24
|
+
locked_ids (set): Track IDs whose speed has been finalized.
|
|
25
|
+
max_hist (int): Required frame history before computing speed.
|
|
26
|
+
meter_per_pixel (float): Real-world meters represented by one pixel for scene scale conversion.
|
|
27
|
+
max_speed (int): Maximum allowed object speed; values above this will be capped.
|
|
28
|
+
|
|
29
|
+
Methods:
|
|
30
|
+
process: Process input frames to estimate object speeds based on tracking data.
|
|
31
|
+
store_tracking_history: Store the tracking history for an object.
|
|
32
|
+
extract_tracks: Extract tracks from the current frame.
|
|
33
|
+
display_output: Display the output with annotations.
|
|
34
|
+
|
|
35
|
+
Examples:
|
|
36
|
+
Initialize speed estimator and process a frame
|
|
37
|
+
>>> estimator = SpeedEstimator(meter_per_pixel=0.04, max_speed=120)
|
|
38
|
+
>>> frame = cv2.imread("frame.jpg")
|
|
39
|
+
>>> results = estimator.process(frame)
|
|
40
|
+
>>> cv2.imshow("Speed Estimation", results.plot_im)
|
|
41
|
+
"""
|
|
42
|
+
|
|
43
|
+
def __init__(self, **kwargs: Any) -> None:
|
|
44
|
+
"""Initialize the SpeedEstimator object with speed estimation parameters and data structures.
|
|
45
|
+
|
|
46
|
+
Args:
|
|
47
|
+
**kwargs (Any): Additional keyword arguments passed to the parent class.
|
|
48
|
+
"""
|
|
49
|
+
super().__init__(**kwargs)
|
|
50
|
+
|
|
51
|
+
self.fps = self.CFG["fps"] # Video frame rate for time calculations
|
|
52
|
+
self.frame_count = 0 # Global frame counter
|
|
53
|
+
self.trk_frame_ids = {} # Track ID → first frame index
|
|
54
|
+
self.spd = {} # Final speed per object (km/h), once locked
|
|
55
|
+
self.trk_hist = {} # Track ID → deque of (time, position)
|
|
56
|
+
self.locked_ids = set() # Track IDs whose speed has been finalized
|
|
57
|
+
self.max_hist = self.CFG["max_hist"] # Required frame history before computing speed
|
|
58
|
+
self.meter_per_pixel = self.CFG["meter_per_pixel"] # Scene scale, depends on camera details
|
|
59
|
+
self.max_speed = self.CFG["max_speed"] # Maximum speed adjustment
|
|
60
|
+
|
|
61
|
+
def process(self, im0) -> SolutionResults:
|
|
62
|
+
"""Process an input frame to estimate object speeds based on tracking data.
|
|
63
|
+
|
|
64
|
+
Args:
|
|
65
|
+
im0 (np.ndarray): Input image for processing with shape (H, W, C) in OpenCV BGR format.
|
|
66
|
+
|
|
67
|
+
Returns:
|
|
68
|
+
(SolutionResults): Contains processed image `plot_im` and `total_tracks` (number of tracked objects).
|
|
69
|
+
|
|
70
|
+
Examples:
|
|
71
|
+
Process a frame for speed estimation
|
|
72
|
+
>>> estimator = SpeedEstimator()
|
|
73
|
+
>>> image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
|
|
74
|
+
>>> results = estimator.process(image)
|
|
75
|
+
"""
|
|
76
|
+
self.frame_count += 1
|
|
77
|
+
self.extract_tracks(im0)
|
|
78
|
+
annotator = SolutionAnnotator(im0, line_width=self.line_width)
|
|
79
|
+
|
|
80
|
+
for box, track_id, _, _ in zip(self.boxes, self.track_ids, self.clss, self.confs):
|
|
81
|
+
self.store_tracking_history(track_id, box)
|
|
82
|
+
|
|
83
|
+
if track_id not in self.trk_hist: # Initialize history if new track found
|
|
84
|
+
self.trk_hist[track_id] = deque(maxlen=self.max_hist)
|
|
85
|
+
self.trk_frame_ids[track_id] = self.frame_count
|
|
86
|
+
|
|
87
|
+
if track_id not in self.locked_ids: # Update history until speed is locked
|
|
88
|
+
trk_hist = self.trk_hist[track_id]
|
|
89
|
+
trk_hist.append(self.track_line[-1])
|
|
90
|
+
|
|
91
|
+
# Compute and lock speed once enough history is collected
|
|
92
|
+
if len(trk_hist) == self.max_hist:
|
|
93
|
+
p0, p1 = trk_hist[0], trk_hist[-1] # First and last points of track
|
|
94
|
+
dt = (self.frame_count - self.trk_frame_ids[track_id]) / self.fps # Time in seconds
|
|
95
|
+
if dt > 0:
|
|
96
|
+
dx, dy = p1[0] - p0[0], p1[1] - p0[1] # Pixel displacement
|
|
97
|
+
pixel_distance = sqrt(dx * dx + dy * dy) # Calculate pixel distance
|
|
98
|
+
meters = pixel_distance * self.meter_per_pixel # Convert to meters
|
|
99
|
+
self.spd[track_id] = int(
|
|
100
|
+
min((meters / dt) * 3.6, self.max_speed)
|
|
101
|
+
) # Convert to km/h and store final speed
|
|
102
|
+
self.locked_ids.add(track_id) # Prevent further updates
|
|
103
|
+
self.trk_hist.pop(track_id, None) # Free memory
|
|
104
|
+
self.trk_frame_ids.pop(track_id, None) # Remove frame start reference
|
|
105
|
+
|
|
106
|
+
if track_id in self.spd:
|
|
107
|
+
speed_label = f"{self.spd[track_id]} km/h"
|
|
108
|
+
annotator.box_label(box, label=speed_label, color=colors(track_id, True)) # Draw bounding box
|
|
109
|
+
|
|
110
|
+
plot_im = annotator.result()
|
|
111
|
+
self.display_output(plot_im) # Display output with base class function
|
|
112
|
+
|
|
113
|
+
# Return results with processed image and tracking summary
|
|
114
|
+
return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids))
|
|
@@ -0,0 +1,260 @@
|
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
2
|
+
|
|
3
|
+
import io
|
|
4
|
+
import os
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
import cv2
|
|
8
|
+
import torch
|
|
9
|
+
|
|
10
|
+
from ultralytics import YOLO
|
|
11
|
+
from ultralytics.utils import LOGGER
|
|
12
|
+
from ultralytics.utils.checks import check_requirements
|
|
13
|
+
from ultralytics.utils.downloads import GITHUB_ASSETS_STEMS
|
|
14
|
+
|
|
15
|
+
torch.classes.__path__ = [] # Torch module __path__._path issue: https://github.com/datalab-to/marker/issues/442
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class Inference:
|
|
19
|
+
"""A class to perform object detection, image classification, image segmentation and pose estimation inference.
|
|
20
|
+
|
|
21
|
+
This class provides functionalities for loading models, configuring settings, uploading video files, and performing
|
|
22
|
+
real-time inference using Streamlit and Ultralytics YOLO models.
|
|
23
|
+
|
|
24
|
+
Attributes:
|
|
25
|
+
st (module): Streamlit module for UI creation.
|
|
26
|
+
temp_dict (dict): Temporary dictionary to store the model path and other configuration.
|
|
27
|
+
model_path (str): Path to the loaded model.
|
|
28
|
+
model (YOLO): The YOLO model instance.
|
|
29
|
+
source (str): Selected video source (webcam or video file).
|
|
30
|
+
enable_trk (bool): Enable tracking option.
|
|
31
|
+
conf (float): Confidence threshold for detection.
|
|
32
|
+
iou (float): IoU threshold for non-maximum suppression.
|
|
33
|
+
org_frame (Any): Container for the original frame to be displayed.
|
|
34
|
+
ann_frame (Any): Container for the annotated frame to be displayed.
|
|
35
|
+
vid_file_name (str | int): Name of the uploaded video file or webcam index.
|
|
36
|
+
selected_ind (list[int]): List of selected class indices for detection.
|
|
37
|
+
|
|
38
|
+
Methods:
|
|
39
|
+
web_ui: Set up the Streamlit web interface with custom HTML elements.
|
|
40
|
+
sidebar: Configure the Streamlit sidebar for model and inference settings.
|
|
41
|
+
source_upload: Handle video file uploads through the Streamlit interface.
|
|
42
|
+
configure: Configure the model and load selected classes for inference.
|
|
43
|
+
inference: Perform real-time object detection inference.
|
|
44
|
+
|
|
45
|
+
Examples:
|
|
46
|
+
Create an Inference instance with a custom model
|
|
47
|
+
>>> inf = Inference(model="path/to/model.pt")
|
|
48
|
+
>>> inf.inference()
|
|
49
|
+
|
|
50
|
+
Create an Inference instance with default settings
|
|
51
|
+
>>> inf = Inference()
|
|
52
|
+
>>> inf.inference()
|
|
53
|
+
"""
|
|
54
|
+
|
|
55
|
+
def __init__(self, **kwargs: Any) -> None:
|
|
56
|
+
"""Initialize the Inference class, checking Streamlit requirements and setting up the model path.
|
|
57
|
+
|
|
58
|
+
Args:
|
|
59
|
+
**kwargs (Any): Additional keyword arguments for model configuration.
|
|
60
|
+
"""
|
|
61
|
+
check_requirements("streamlit>=1.29.0") # scope imports for faster ultralytics package load speeds
|
|
62
|
+
import streamlit as st
|
|
63
|
+
|
|
64
|
+
self.st = st # Reference to the Streamlit module
|
|
65
|
+
self.source = None # Video source selection (webcam or video file)
|
|
66
|
+
self.img_file_names = [] # List of image file names
|
|
67
|
+
self.enable_trk = False # Flag to toggle object tracking
|
|
68
|
+
self.conf = 0.25 # Confidence threshold for detection
|
|
69
|
+
self.iou = 0.45 # Intersection-over-Union (IoU) threshold for non-maximum suppression
|
|
70
|
+
self.org_frame = None # Container for the original frame display
|
|
71
|
+
self.ann_frame = None # Container for the annotated frame display
|
|
72
|
+
self.vid_file_name = None # Video file name or webcam index
|
|
73
|
+
self.selected_ind: list[int] = [] # List of selected class indices for detection
|
|
74
|
+
self.model = None # YOLO model instance
|
|
75
|
+
|
|
76
|
+
self.temp_dict = {"model": None, **kwargs}
|
|
77
|
+
self.model_path = None # Model file path
|
|
78
|
+
if self.temp_dict["model"] is not None:
|
|
79
|
+
self.model_path = self.temp_dict["model"]
|
|
80
|
+
|
|
81
|
+
LOGGER.info(f"Ultralytics Solutions: ✅ {self.temp_dict}")
|
|
82
|
+
|
|
83
|
+
def web_ui(self) -> None:
|
|
84
|
+
"""Set up the Streamlit web interface with custom HTML elements."""
|
|
85
|
+
menu_style_cfg = """<style>MainMenu {visibility: hidden;}</style>""" # Hide main menu style
|
|
86
|
+
|
|
87
|
+
# Main title of streamlit application
|
|
88
|
+
main_title_cfg = """<div><h1 style="color:#111F68; text-align:center; font-size:40px; margin-top:-50px;
|
|
89
|
+
font-family: 'Archivo', sans-serif; margin-bottom:20px;">Ultralytics YOLO Streamlit Application</h1></div>"""
|
|
90
|
+
|
|
91
|
+
# Subtitle of streamlit application
|
|
92
|
+
sub_title_cfg = """<div><h5 style="color:#042AFF; text-align:center; font-family: 'Archivo', sans-serif;
|
|
93
|
+
margin-top:-15px; margin-bottom:50px;">Experience real-time object detection on your webcam, videos, and images
|
|
94
|
+
with the power of Ultralytics YOLO! 🚀</h5></div>"""
|
|
95
|
+
|
|
96
|
+
# Set html page configuration and append custom HTML
|
|
97
|
+
self.st.set_page_config(page_title="Ultralytics Streamlit App", layout="wide")
|
|
98
|
+
self.st.markdown(menu_style_cfg, unsafe_allow_html=True)
|
|
99
|
+
self.st.markdown(main_title_cfg, unsafe_allow_html=True)
|
|
100
|
+
self.st.markdown(sub_title_cfg, unsafe_allow_html=True)
|
|
101
|
+
|
|
102
|
+
def sidebar(self) -> None:
|
|
103
|
+
"""Configure the Streamlit sidebar for model and inference settings."""
|
|
104
|
+
with self.st.sidebar: # Add Ultralytics LOGO
|
|
105
|
+
logo = "https://raw.githubusercontent.com/ultralytics/assets/main/logo/Ultralytics_Logotype_Original.svg"
|
|
106
|
+
self.st.image(logo, width=250)
|
|
107
|
+
|
|
108
|
+
self.st.sidebar.title("User Configuration") # Add elements to vertical setting menu
|
|
109
|
+
self.source = self.st.sidebar.selectbox(
|
|
110
|
+
"Source",
|
|
111
|
+
("webcam", "video", "image"),
|
|
112
|
+
) # Add source selection dropdown
|
|
113
|
+
if self.source in ["webcam", "video"]:
|
|
114
|
+
self.enable_trk = self.st.sidebar.radio("Enable Tracking", ("Yes", "No")) == "Yes" # Enable object tracking
|
|
115
|
+
self.conf = float(
|
|
116
|
+
self.st.sidebar.slider("Confidence Threshold", 0.0, 1.0, self.conf, 0.01)
|
|
117
|
+
) # Slider for confidence
|
|
118
|
+
self.iou = float(self.st.sidebar.slider("IoU Threshold", 0.0, 1.0, self.iou, 0.01)) # Slider for NMS threshold
|
|
119
|
+
|
|
120
|
+
if self.source != "image": # Only create columns for video/webcam
|
|
121
|
+
col1, col2 = self.st.columns(2) # Create two columns for displaying frames
|
|
122
|
+
self.org_frame = col1.empty() # Container for original frame
|
|
123
|
+
self.ann_frame = col2.empty() # Container for annotated frame
|
|
124
|
+
|
|
125
|
+
def source_upload(self) -> None:
|
|
126
|
+
"""Handle video file uploads through the Streamlit interface."""
|
|
127
|
+
from ultralytics.data.utils import IMG_FORMATS, VID_FORMATS # scope import
|
|
128
|
+
|
|
129
|
+
self.vid_file_name = ""
|
|
130
|
+
if self.source == "video":
|
|
131
|
+
vid_file = self.st.sidebar.file_uploader("Upload Video File", type=VID_FORMATS)
|
|
132
|
+
if vid_file is not None:
|
|
133
|
+
g = io.BytesIO(vid_file.read()) # BytesIO Object
|
|
134
|
+
with open("ultralytics.mp4", "wb") as out: # Open temporary file as bytes
|
|
135
|
+
out.write(g.read()) # Read bytes into file
|
|
136
|
+
self.vid_file_name = "ultralytics.mp4"
|
|
137
|
+
elif self.source == "webcam":
|
|
138
|
+
self.vid_file_name = 0 # Use webcam index 0
|
|
139
|
+
elif self.source == "image":
|
|
140
|
+
import tempfile # scope import
|
|
141
|
+
|
|
142
|
+
if imgfiles := self.st.sidebar.file_uploader(
|
|
143
|
+
"Upload Image Files", type=IMG_FORMATS, accept_multiple_files=True
|
|
144
|
+
):
|
|
145
|
+
for imgfile in imgfiles: # Save each uploaded image to a temporary file
|
|
146
|
+
with tempfile.NamedTemporaryFile(delete=False, suffix=f".{imgfile.name.split('.')[-1]}") as tf:
|
|
147
|
+
tf.write(imgfile.read())
|
|
148
|
+
self.img_file_names.append({"path": tf.name, "name": imgfile.name})
|
|
149
|
+
|
|
150
|
+
def configure(self) -> None:
|
|
151
|
+
"""Configure the model and load selected classes for inference."""
|
|
152
|
+
# Add dropdown menu for model selection
|
|
153
|
+
M_ORD, T_ORD = ["yolo11n", "yolo11s", "yolo11m", "yolo11l", "yolo11x"], ["", "-seg", "-pose", "-obb", "-cls"]
|
|
154
|
+
available_models = sorted(
|
|
155
|
+
[
|
|
156
|
+
x.replace("yolo", "YOLO")
|
|
157
|
+
for x in GITHUB_ASSETS_STEMS
|
|
158
|
+
if any(x.startswith(b) for b in M_ORD) and "grayscale" not in x
|
|
159
|
+
],
|
|
160
|
+
key=lambda x: (M_ORD.index(x[:7].lower()), T_ORD.index(x[7:].lower() or "")),
|
|
161
|
+
)
|
|
162
|
+
if self.model_path: # Insert user provided custom model in available_models
|
|
163
|
+
available_models.insert(0, self.model_path)
|
|
164
|
+
selected_model = self.st.sidebar.selectbox("Model", available_models)
|
|
165
|
+
|
|
166
|
+
with self.st.spinner("Model is downloading..."):
|
|
167
|
+
if selected_model.endswith((".pt", ".onnx", ".torchscript", ".mlpackage", ".engine")) or any(
|
|
168
|
+
fmt in selected_model for fmt in ("openvino_model", "rknn_model")
|
|
169
|
+
):
|
|
170
|
+
model_path = selected_model
|
|
171
|
+
else:
|
|
172
|
+
model_path = f"{selected_model.lower()}.pt" # Default to .pt if no model provided during function call.
|
|
173
|
+
self.model = YOLO(model_path) # Load the YOLO model
|
|
174
|
+
class_names = list(self.model.names.values()) # Convert dictionary to list of class names
|
|
175
|
+
self.st.success("Model loaded successfully!")
|
|
176
|
+
|
|
177
|
+
# Multiselect box with class names and get indices of selected classes
|
|
178
|
+
selected_classes = self.st.sidebar.multiselect("Classes", class_names, default=class_names[:3])
|
|
179
|
+
self.selected_ind = [class_names.index(option) for option in selected_classes]
|
|
180
|
+
|
|
181
|
+
if not isinstance(self.selected_ind, list): # Ensure selected_options is a list
|
|
182
|
+
self.selected_ind = list(self.selected_ind)
|
|
183
|
+
|
|
184
|
+
def image_inference(self) -> None:
|
|
185
|
+
"""Perform inference on uploaded images."""
|
|
186
|
+
for img_info in self.img_file_names:
|
|
187
|
+
img_path = img_info["path"]
|
|
188
|
+
image = cv2.imread(img_path) # Load and display the original image
|
|
189
|
+
if image is not None:
|
|
190
|
+
self.st.markdown(f"#### Processed: {img_info['name']}")
|
|
191
|
+
col1, col2 = self.st.columns(2)
|
|
192
|
+
with col1:
|
|
193
|
+
self.st.image(image, channels="BGR", caption="Original Image")
|
|
194
|
+
results = self.model(image, conf=self.conf, iou=self.iou, classes=self.selected_ind)
|
|
195
|
+
annotated_image = results[0].plot()
|
|
196
|
+
with col2:
|
|
197
|
+
self.st.image(annotated_image, channels="BGR", caption="Predicted Image")
|
|
198
|
+
try: # Clean up temporary file
|
|
199
|
+
os.unlink(img_path)
|
|
200
|
+
except FileNotFoundError:
|
|
201
|
+
pass # File doesn't exist, ignore
|
|
202
|
+
else:
|
|
203
|
+
self.st.error("Could not load the uploaded image.")
|
|
204
|
+
|
|
205
|
+
def inference(self) -> None:
|
|
206
|
+
"""Perform real-time object detection inference on video or webcam feed."""
|
|
207
|
+
self.web_ui() # Initialize the web interface
|
|
208
|
+
self.sidebar() # Create the sidebar
|
|
209
|
+
self.source_upload() # Upload the video source
|
|
210
|
+
self.configure() # Configure the app
|
|
211
|
+
|
|
212
|
+
if self.st.sidebar.button("Start"):
|
|
213
|
+
if self.source == "image":
|
|
214
|
+
if self.img_file_names:
|
|
215
|
+
self.image_inference()
|
|
216
|
+
else:
|
|
217
|
+
self.st.info("Please upload an image file to perform inference.")
|
|
218
|
+
return
|
|
219
|
+
|
|
220
|
+
stop_button = self.st.sidebar.button("Stop") # Button to stop the inference
|
|
221
|
+
cap = cv2.VideoCapture(self.vid_file_name) # Capture the video
|
|
222
|
+
if not cap.isOpened():
|
|
223
|
+
self.st.error("Could not open webcam or video source.")
|
|
224
|
+
return
|
|
225
|
+
|
|
226
|
+
while cap.isOpened():
|
|
227
|
+
success, frame = cap.read()
|
|
228
|
+
if not success:
|
|
229
|
+
self.st.warning("Failed to read frame from webcam. Please verify the webcam is connected properly.")
|
|
230
|
+
break
|
|
231
|
+
|
|
232
|
+
# Process frame with model
|
|
233
|
+
if self.enable_trk:
|
|
234
|
+
results = self.model.track(
|
|
235
|
+
frame, conf=self.conf, iou=self.iou, classes=self.selected_ind, persist=True
|
|
236
|
+
)
|
|
237
|
+
else:
|
|
238
|
+
results = self.model(frame, conf=self.conf, iou=self.iou, classes=self.selected_ind)
|
|
239
|
+
|
|
240
|
+
annotated_frame = results[0].plot() # Add annotations on frame
|
|
241
|
+
|
|
242
|
+
if stop_button:
|
|
243
|
+
cap.release() # Release the capture
|
|
244
|
+
self.st.stop() # Stop streamlit app
|
|
245
|
+
|
|
246
|
+
self.org_frame.image(frame, channels="BGR", caption="Original Frame") # Display original frame
|
|
247
|
+
self.ann_frame.image(annotated_frame, channels="BGR", caption="Predicted Frame") # Display processed
|
|
248
|
+
|
|
249
|
+
cap.release() # Release the capture
|
|
250
|
+
cv2.destroyAllWindows() # Destroy all OpenCV windows
|
|
251
|
+
|
|
252
|
+
|
|
253
|
+
if __name__ == "__main__":
|
|
254
|
+
import sys # Import the sys module for accessing command-line arguments
|
|
255
|
+
|
|
256
|
+
# Check if a model name is provided as a command-line argument
|
|
257
|
+
args = len(sys.argv)
|
|
258
|
+
model = sys.argv[1] if args > 1 else None # Assign first argument as the model name if provided
|
|
259
|
+
# Create an instance of the Inference class and run inference
|
|
260
|
+
Inference(model=model).inference()
|
|
@@ -0,0 +1,156 @@
|
|
|
1
|
+
<!-- Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license -->
|
|
2
|
+
|
|
3
|
+
<!--Similarity search webpage-->
|
|
4
|
+
<!doctype html>
|
|
5
|
+
<html lang="en">
|
|
6
|
+
<head>
|
|
7
|
+
<meta charset="UTF-8" />
|
|
8
|
+
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
|
9
|
+
<title>Semantic Image Search</title>
|
|
10
|
+
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap" rel="stylesheet" />
|
|
11
|
+
<style>
|
|
12
|
+
body {
|
|
13
|
+
background: linear-gradient(135deg, #f0f4ff, #f9fbff);
|
|
14
|
+
font-family: "Inter", sans-serif;
|
|
15
|
+
color: #111e68;
|
|
16
|
+
padding: 2rem;
|
|
17
|
+
margin: 0;
|
|
18
|
+
min-height: 100vh;
|
|
19
|
+
}
|
|
20
|
+
|
|
21
|
+
h1 {
|
|
22
|
+
text-align: center;
|
|
23
|
+
margin-bottom: 2rem;
|
|
24
|
+
font-size: 2.5rem;
|
|
25
|
+
font-weight: 600;
|
|
26
|
+
}
|
|
27
|
+
|
|
28
|
+
form {
|
|
29
|
+
display: flex;
|
|
30
|
+
flex-wrap: wrap;
|
|
31
|
+
justify-content: center;
|
|
32
|
+
align-items: center;
|
|
33
|
+
gap: 1rem;
|
|
34
|
+
margin-bottom: 3rem;
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
input[type="text"] {
|
|
38
|
+
width: 300px;
|
|
39
|
+
padding: 0.75rem 1rem;
|
|
40
|
+
font-size: 1rem;
|
|
41
|
+
border-radius: 10px;
|
|
42
|
+
border: 1px solid #ccc;
|
|
43
|
+
box-shadow: 0 2px 6px rgba(0, 0, 0, 0.05);
|
|
44
|
+
transition: box-shadow 0.3s ease;
|
|
45
|
+
}
|
|
46
|
+
|
|
47
|
+
input[type="text"]:focus {
|
|
48
|
+
outline: none;
|
|
49
|
+
box-shadow: 0 0 0 3px rgba(17, 30, 104, 0.2);
|
|
50
|
+
}
|
|
51
|
+
|
|
52
|
+
button {
|
|
53
|
+
background-color: #111e68;
|
|
54
|
+
color: white;
|
|
55
|
+
font-weight: 600;
|
|
56
|
+
font-size: 1rem;
|
|
57
|
+
padding: 0.75rem 1.5rem;
|
|
58
|
+
border-radius: 10px;
|
|
59
|
+
border: none;
|
|
60
|
+
cursor: pointer;
|
|
61
|
+
transition:
|
|
62
|
+
background-color 0.3s ease,
|
|
63
|
+
transform 0.2s ease;
|
|
64
|
+
}
|
|
65
|
+
|
|
66
|
+
button:hover {
|
|
67
|
+
background-color: #1f2e9f;
|
|
68
|
+
transform: translateY(-2px);
|
|
69
|
+
}
|
|
70
|
+
|
|
71
|
+
.grid {
|
|
72
|
+
display: grid;
|
|
73
|
+
grid-template-columns: repeat(auto-fill, minmax(260px, 1fr));
|
|
74
|
+
gap: 1.5rem;
|
|
75
|
+
max-width: 1600px;
|
|
76
|
+
margin: auto;
|
|
77
|
+
}
|
|
78
|
+
|
|
79
|
+
.card {
|
|
80
|
+
background: white;
|
|
81
|
+
border-radius: 16px;
|
|
82
|
+
overflow: hidden;
|
|
83
|
+
box-shadow: 0 6px 14px rgba(0, 0, 0, 0.08);
|
|
84
|
+
transition:
|
|
85
|
+
transform 0.3s ease,
|
|
86
|
+
box-shadow 0.3s ease;
|
|
87
|
+
}
|
|
88
|
+
|
|
89
|
+
.card:hover {
|
|
90
|
+
transform: translateY(-6px);
|
|
91
|
+
box-shadow: 0 10px 20px rgba(0, 0, 0, 0.1);
|
|
92
|
+
}
|
|
93
|
+
|
|
94
|
+
.card img {
|
|
95
|
+
width: 100%;
|
|
96
|
+
height: 100%;
|
|
97
|
+
object-fit: cover;
|
|
98
|
+
display: block;
|
|
99
|
+
}
|
|
100
|
+
</style>
|
|
101
|
+
</head>
|
|
102
|
+
<script>
|
|
103
|
+
function filterResults(k) {
|
|
104
|
+
const cards = document.querySelectorAll(".grid .card");
|
|
105
|
+
cards.forEach((card, idx) => {
|
|
106
|
+
card.style.display = idx < k ? "block" : "none";
|
|
107
|
+
});
|
|
108
|
+
const buttons = document.querySelectorAll(".topk-btn");
|
|
109
|
+
buttons.forEach((btn) => {
|
|
110
|
+
btn.classList.remove("active");
|
|
111
|
+
});
|
|
112
|
+
event.target.classList.add("active");
|
|
113
|
+
}
|
|
114
|
+
document.addEventListener("DOMContentLoaded", () => {
|
|
115
|
+
filterResults(10);
|
|
116
|
+
});
|
|
117
|
+
</script>
|
|
118
|
+
<body>
|
|
119
|
+
<div style="text-align: center; margin-bottom: 1rem">
|
|
120
|
+
<img
|
|
121
|
+
src="https://raw.githubusercontent.com/ultralytics/assets/main/logo/favicon.png"
|
|
122
|
+
alt="Ultralytics Logo"
|
|
123
|
+
style="height: 40px"
|
|
124
|
+
/>
|
|
125
|
+
</div>
|
|
126
|
+
<h1>Semantic Image Search with AI</h1>
|
|
127
|
+
|
|
128
|
+
<!-- Search box -->
|
|
129
|
+
<form method="POST">
|
|
130
|
+
<input
|
|
131
|
+
type="text"
|
|
132
|
+
name="query"
|
|
133
|
+
placeholder="Describe the scene (e.g., man walking)"
|
|
134
|
+
value="{{ request.form['query'] }}"
|
|
135
|
+
required
|
|
136
|
+
/>
|
|
137
|
+
<button type="submit">Search</button>
|
|
138
|
+
{% if results %}
|
|
139
|
+
<div class="top-k-buttons">
|
|
140
|
+
<button type="button" class="topk-btn" onclick="filterResults(5)">Top 5</button>
|
|
141
|
+
<button type="button" class="topk-btn active" onclick="filterResults(10)">Top 10</button>
|
|
142
|
+
<button type="button" class="topk-btn" onclick="filterResults(30)">Top 30</button>
|
|
143
|
+
</div>
|
|
144
|
+
{% endif %}
|
|
145
|
+
</form>
|
|
146
|
+
|
|
147
|
+
<!-- Search results grid -->
|
|
148
|
+
<div class="grid">
|
|
149
|
+
{% for img in results %}
|
|
150
|
+
<div class="card">
|
|
151
|
+
<img src="{{ url_for('static', filename=img) }}" alt="Result Image" />
|
|
152
|
+
</div>
|
|
153
|
+
{% endfor %}
|
|
154
|
+
</div>
|
|
155
|
+
</body>
|
|
156
|
+
</html>
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
2
|
+
|
|
3
|
+
from typing import Any
|
|
4
|
+
|
|
5
|
+
import cv2
|
|
6
|
+
import numpy as np
|
|
7
|
+
|
|
8
|
+
from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
|
|
9
|
+
from ultralytics.utils.plotting import colors
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class TrackZone(BaseSolution):
|
|
13
|
+
"""A class to manage region-based object tracking in a video stream.
|
|
14
|
+
|
|
15
|
+
This class extends the BaseSolution class and provides functionality for tracking objects within a specific region
|
|
16
|
+
defined by a polygonal area. Objects outside the region are excluded from tracking.
|
|
17
|
+
|
|
18
|
+
Attributes:
|
|
19
|
+
region (np.ndarray): The polygonal region for tracking, represented as a convex hull of points.
|
|
20
|
+
line_width (int): Width of the lines used for drawing bounding boxes and region boundaries.
|
|
21
|
+
names (list[str]): List of class names that the model can detect.
|
|
22
|
+
boxes (list[np.ndarray]): Bounding boxes of tracked objects.
|
|
23
|
+
track_ids (list[int]): Unique identifiers for each tracked object.
|
|
24
|
+
clss (list[int]): Class indices of tracked objects.
|
|
25
|
+
|
|
26
|
+
Methods:
|
|
27
|
+
process: Process each frame of the video, applying region-based tracking.
|
|
28
|
+
extract_tracks: Extract tracking information from the input frame.
|
|
29
|
+
display_output: Display the processed output.
|
|
30
|
+
|
|
31
|
+
Examples:
|
|
32
|
+
>>> tracker = TrackZone()
|
|
33
|
+
>>> frame = cv2.imread("frame.jpg")
|
|
34
|
+
>>> results = tracker.process(frame)
|
|
35
|
+
>>> cv2.imshow("Tracked Frame", results.plot_im)
|
|
36
|
+
"""
|
|
37
|
+
|
|
38
|
+
def __init__(self, **kwargs: Any) -> None:
|
|
39
|
+
"""Initialize the TrackZone class for tracking objects within a defined region in video streams.
|
|
40
|
+
|
|
41
|
+
Args:
|
|
42
|
+
**kwargs (Any): Additional keyword arguments passed to the parent class.
|
|
43
|
+
"""
|
|
44
|
+
super().__init__(**kwargs)
|
|
45
|
+
default_region = [(75, 75), (565, 75), (565, 285), (75, 285)]
|
|
46
|
+
self.region = cv2.convexHull(np.array(self.region or default_region, dtype=np.int32))
|
|
47
|
+
self.mask = None
|
|
48
|
+
|
|
49
|
+
def process(self, im0: np.ndarray) -> SolutionResults:
|
|
50
|
+
"""Process the input frame to track objects within a defined region.
|
|
51
|
+
|
|
52
|
+
This method initializes the annotator, creates a mask for the specified region, extracts tracks only from the
|
|
53
|
+
masked area, and updates tracking information. Objects outside the region are ignored.
|
|
54
|
+
|
|
55
|
+
Args:
|
|
56
|
+
im0 (np.ndarray): The input image or frame to be processed.
|
|
57
|
+
|
|
58
|
+
Returns:
|
|
59
|
+
(SolutionResults): Contains processed image `plot_im` and `total_tracks` (int) representing the total number
|
|
60
|
+
of tracked objects within the defined region.
|
|
61
|
+
|
|
62
|
+
Examples:
|
|
63
|
+
>>> tracker = TrackZone()
|
|
64
|
+
>>> frame = cv2.imread("path/to/image.jpg")
|
|
65
|
+
>>> results = tracker.process(frame)
|
|
66
|
+
"""
|
|
67
|
+
annotator = SolutionAnnotator(im0, line_width=self.line_width) # Initialize annotator
|
|
68
|
+
|
|
69
|
+
if self.mask is None: # Create a mask for the region
|
|
70
|
+
self.mask = np.zeros_like(im0[:, :, 0])
|
|
71
|
+
cv2.fillPoly(self.mask, [self.region], 255)
|
|
72
|
+
masked_frame = cv2.bitwise_and(im0, im0, mask=self.mask)
|
|
73
|
+
self.extract_tracks(masked_frame)
|
|
74
|
+
|
|
75
|
+
# Draw the region boundary
|
|
76
|
+
cv2.polylines(im0, [self.region], isClosed=True, color=(255, 255, 255), thickness=self.line_width * 2)
|
|
77
|
+
|
|
78
|
+
# Iterate over boxes, track ids, classes indexes list and draw bounding boxes
|
|
79
|
+
for box, track_id, cls, conf in zip(self.boxes, self.track_ids, self.clss, self.confs):
|
|
80
|
+
annotator.box_label(
|
|
81
|
+
box, label=self.adjust_box_label(cls, conf, track_id=track_id), color=colors(track_id, True)
|
|
82
|
+
)
|
|
83
|
+
|
|
84
|
+
plot_im = annotator.result()
|
|
85
|
+
self.display_output(plot_im) # Display output with base class function
|
|
86
|
+
|
|
87
|
+
# Return a SolutionResults
|
|
88
|
+
return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids))
|