dgenerate-ultralytics-headless 8.3.253__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.
- dgenerate_ultralytics_headless-8.3.253.dist-info/METADATA +405 -0
- dgenerate_ultralytics_headless-8.3.253.dist-info/RECORD +299 -0
- dgenerate_ultralytics_headless-8.3.253.dist-info/WHEEL +5 -0
- dgenerate_ultralytics_headless-8.3.253.dist-info/entry_points.txt +3 -0
- dgenerate_ultralytics_headless-8.3.253.dist-info/licenses/LICENSE +661 -0
- dgenerate_ultralytics_headless-8.3.253.dist-info/top_level.txt +1 -0
- 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 +1028 -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/TT100K.yaml +346 -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 +1580 -0
- ultralytics/engine/model.py +1125 -0
- ultralytics/engine/predictor.py +508 -0
- ultralytics/engine/results.py +1522 -0
- ultralytics/engine/trainer.py +977 -0
- ultralytics/engine/tuner.py +449 -0
- ultralytics/engine/validator.py +387 -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 +964 -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 +453 -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 +1020 -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 +325 -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 +506 -0
- ultralytics/utils/loss.py +849 -0
- ultralytics/utils/metrics.py +1563 -0
- ultralytics/utils/nms.py +337 -0
- ultralytics/utils/ops.py +664 -0
- ultralytics/utils/patches.py +201 -0
- ultralytics/utils/plotting.py +1047 -0
- ultralytics/utils/tal.py +404 -0
- ultralytics/utils/torch_utils.py +984 -0
- ultralytics/utils/tqdm.py +443 -0
- ultralytics/utils/triton.py +112 -0
- ultralytics/utils/tuner.py +168 -0
|
@@ -0,0 +1,231 @@
|
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
|
|
7
|
+
import numpy as np
|
|
8
|
+
import torch
|
|
9
|
+
|
|
10
|
+
from ultralytics.nn.modules import Detect, Pose
|
|
11
|
+
from ultralytics.utils import LOGGER
|
|
12
|
+
from ultralytics.utils.downloads import attempt_download_asset
|
|
13
|
+
from ultralytics.utils.files import spaces_in_path
|
|
14
|
+
from ultralytics.utils.tal import make_anchors
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def tf_wrapper(model: torch.nn.Module) -> torch.nn.Module:
|
|
18
|
+
"""A wrapper to add TensorFlow compatible inference methods to Detect and Pose layers."""
|
|
19
|
+
for m in model.modules():
|
|
20
|
+
if not isinstance(m, Detect):
|
|
21
|
+
continue
|
|
22
|
+
import types
|
|
23
|
+
|
|
24
|
+
m._inference = types.MethodType(_tf_inference, m)
|
|
25
|
+
if type(m) is Pose:
|
|
26
|
+
m.kpts_decode = types.MethodType(tf_kpts_decode, m)
|
|
27
|
+
return model
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _tf_inference(self, x: list[torch.Tensor]) -> tuple[torch.Tensor]:
|
|
31
|
+
"""Decode boxes and cls scores for tf object detection."""
|
|
32
|
+
shape = x[0].shape # BCHW
|
|
33
|
+
x_cat = torch.cat([xi.view(x[0].shape[0], self.no, -1) for xi in x], 2)
|
|
34
|
+
box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
|
|
35
|
+
if self.dynamic or self.shape != shape:
|
|
36
|
+
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
|
|
37
|
+
self.shape = shape
|
|
38
|
+
grid_h, grid_w = shape[2], shape[3]
|
|
39
|
+
grid_size = torch.tensor([grid_w, grid_h, grid_w, grid_h], device=box.device).reshape(1, 4, 1)
|
|
40
|
+
norm = self.strides / (self.stride[0] * grid_size)
|
|
41
|
+
dbox = self.decode_bboxes(self.dfl(box) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
|
|
42
|
+
return torch.cat((dbox, cls.sigmoid()), 1)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def tf_kpts_decode(self, bs: int, kpts: torch.Tensor) -> torch.Tensor:
|
|
46
|
+
"""Decode keypoints for tf pose estimation."""
|
|
47
|
+
ndim = self.kpt_shape[1]
|
|
48
|
+
# required for TFLite export to avoid 'PLACEHOLDER_FOR_GREATER_OP_CODES' bug
|
|
49
|
+
# Precompute normalization factor to increase numerical stability
|
|
50
|
+
y = kpts.view(bs, *self.kpt_shape, -1)
|
|
51
|
+
grid_h, grid_w = self.shape[2], self.shape[3]
|
|
52
|
+
grid_size = torch.tensor([grid_w, grid_h], device=y.device).reshape(1, 2, 1)
|
|
53
|
+
norm = self.strides / (self.stride[0] * grid_size)
|
|
54
|
+
a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * norm
|
|
55
|
+
if ndim == 3:
|
|
56
|
+
a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
|
|
57
|
+
return a.view(bs, self.nk, -1)
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def onnx2saved_model(
|
|
61
|
+
onnx_file: str,
|
|
62
|
+
output_dir: Path,
|
|
63
|
+
int8: bool = False,
|
|
64
|
+
images: np.ndarray = None,
|
|
65
|
+
disable_group_convolution: bool = False,
|
|
66
|
+
prefix="",
|
|
67
|
+
):
|
|
68
|
+
"""Convert a ONNX model to TensorFlow SavedModel format via ONNX.
|
|
69
|
+
|
|
70
|
+
Args:
|
|
71
|
+
onnx_file (str): ONNX file path.
|
|
72
|
+
output_dir (Path): Output directory path for the SavedModel.
|
|
73
|
+
int8 (bool, optional): Enable INT8 quantization. Defaults to False.
|
|
74
|
+
images (np.ndarray, optional): Calibration images for INT8 quantization in BHWC format.
|
|
75
|
+
disable_group_convolution (bool, optional): Disable group convolution optimization. Defaults to False.
|
|
76
|
+
prefix (str, optional): Logging prefix. Defaults to "".
|
|
77
|
+
|
|
78
|
+
Returns:
|
|
79
|
+
(keras.Model): Converted Keras model.
|
|
80
|
+
|
|
81
|
+
Notes:
|
|
82
|
+
- Requires onnx2tf package. Downloads calibration data if INT8 quantization is enabled.
|
|
83
|
+
- Removes temporary files and renames quantized models after conversion.
|
|
84
|
+
"""
|
|
85
|
+
# Pre-download calibration file to fix https://github.com/PINTO0309/onnx2tf/issues/545
|
|
86
|
+
onnx2tf_file = Path("calibration_image_sample_data_20x128x128x3_float32.npy")
|
|
87
|
+
if not onnx2tf_file.exists():
|
|
88
|
+
attempt_download_asset(f"{onnx2tf_file}.zip", unzip=True, delete=True)
|
|
89
|
+
np_data = None
|
|
90
|
+
if int8:
|
|
91
|
+
tmp_file = output_dir / "tmp_tflite_int8_calibration_images.npy" # int8 calibration images file
|
|
92
|
+
if images is not None:
|
|
93
|
+
output_dir.mkdir(parents=True, exist_ok=True)
|
|
94
|
+
np.save(str(tmp_file), images) # BHWC
|
|
95
|
+
np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
|
|
96
|
+
|
|
97
|
+
# Patch onnx.helper for onnx_graphsurgeon compatibility with ONNX>=1.17
|
|
98
|
+
# The float32_to_bfloat16 function was removed in ONNX 1.17, but onnx_graphsurgeon still uses it
|
|
99
|
+
import onnx.helper
|
|
100
|
+
|
|
101
|
+
if not hasattr(onnx.helper, "float32_to_bfloat16"):
|
|
102
|
+
import struct
|
|
103
|
+
|
|
104
|
+
def float32_to_bfloat16(fval):
|
|
105
|
+
"""Convert float32 to bfloat16 (truncates lower 16 bits of mantissa)."""
|
|
106
|
+
ival = struct.unpack("=I", struct.pack("=f", fval))[0]
|
|
107
|
+
return ival >> 16
|
|
108
|
+
|
|
109
|
+
onnx.helper.float32_to_bfloat16 = float32_to_bfloat16
|
|
110
|
+
|
|
111
|
+
import onnx2tf # scoped for after ONNX export for reduced conflict during import
|
|
112
|
+
|
|
113
|
+
LOGGER.info(f"{prefix} starting TFLite export with onnx2tf {onnx2tf.__version__}...")
|
|
114
|
+
keras_model = onnx2tf.convert(
|
|
115
|
+
input_onnx_file_path=onnx_file,
|
|
116
|
+
output_folder_path=str(output_dir),
|
|
117
|
+
not_use_onnxsim=True,
|
|
118
|
+
verbosity="error", # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
|
|
119
|
+
output_integer_quantized_tflite=int8,
|
|
120
|
+
custom_input_op_name_np_data_path=np_data,
|
|
121
|
+
enable_batchmatmul_unfold=True and not int8, # fix lower no. of detected objects on GPU delegate
|
|
122
|
+
output_signaturedefs=True, # fix error with Attention block group convolution
|
|
123
|
+
disable_group_convolution=disable_group_convolution, # fix error with group convolution
|
|
124
|
+
)
|
|
125
|
+
|
|
126
|
+
# Remove/rename TFLite models
|
|
127
|
+
if int8:
|
|
128
|
+
tmp_file.unlink(missing_ok=True)
|
|
129
|
+
for file in output_dir.rglob("*_dynamic_range_quant.tflite"):
|
|
130
|
+
file.rename(file.with_name(file.stem.replace("_dynamic_range_quant", "_int8") + file.suffix))
|
|
131
|
+
for file in output_dir.rglob("*_integer_quant_with_int16_act.tflite"):
|
|
132
|
+
file.unlink() # delete extra fp16 activation TFLite files
|
|
133
|
+
return keras_model
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
def keras2pb(keras_model, file: Path, prefix=""):
|
|
137
|
+
"""Convert a Keras model to TensorFlow GraphDef (.pb) format.
|
|
138
|
+
|
|
139
|
+
Args:
|
|
140
|
+
keras_model (keras.Model): Keras model to convert to frozen graph format.
|
|
141
|
+
file (Path): Output file path (suffix will be changed to .pb).
|
|
142
|
+
prefix (str, optional): Logging prefix. Defaults to "".
|
|
143
|
+
|
|
144
|
+
Notes:
|
|
145
|
+
Creates a frozen graph by converting variables to constants for inference optimization.
|
|
146
|
+
"""
|
|
147
|
+
import tensorflow as tf
|
|
148
|
+
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
|
|
149
|
+
|
|
150
|
+
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
|
|
151
|
+
m = tf.function(lambda x: keras_model(x)) # full model
|
|
152
|
+
m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
|
|
153
|
+
frozen_func = convert_variables_to_constants_v2(m)
|
|
154
|
+
frozen_func.graph.as_graph_def()
|
|
155
|
+
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(file.parent), name=file.name, as_text=False)
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
def tflite2edgetpu(tflite_file: str | Path, output_dir: str | Path, prefix: str = ""):
|
|
159
|
+
"""Convert a TensorFlow Lite model to Edge TPU format using the Edge TPU compiler.
|
|
160
|
+
|
|
161
|
+
Args:
|
|
162
|
+
tflite_file (str | Path): Path to the input TensorFlow Lite (.tflite) model file.
|
|
163
|
+
output_dir (str | Path): Output directory path for the compiled Edge TPU model.
|
|
164
|
+
prefix (str, optional): Logging prefix. Defaults to "".
|
|
165
|
+
|
|
166
|
+
Notes:
|
|
167
|
+
Requires the Edge TPU compiler to be installed. The function compiles the TFLite model
|
|
168
|
+
for optimal performance on Google's Edge TPU hardware accelerator.
|
|
169
|
+
"""
|
|
170
|
+
import subprocess
|
|
171
|
+
|
|
172
|
+
cmd = (
|
|
173
|
+
"edgetpu_compiler "
|
|
174
|
+
f'--out_dir "{output_dir}" '
|
|
175
|
+
"--show_operations "
|
|
176
|
+
"--search_delegate "
|
|
177
|
+
"--delegate_search_step 30 "
|
|
178
|
+
"--timeout_sec 180 "
|
|
179
|
+
f'"{tflite_file}"'
|
|
180
|
+
)
|
|
181
|
+
LOGGER.info(f"{prefix} running '{cmd}'")
|
|
182
|
+
subprocess.run(cmd, shell=True)
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = False, prefix: str = ""):
|
|
186
|
+
"""Convert a TensorFlow GraphDef (.pb) model to TensorFlow.js format.
|
|
187
|
+
|
|
188
|
+
Args:
|
|
189
|
+
pb_file (str): Path to the input TensorFlow GraphDef (.pb) model file.
|
|
190
|
+
output_dir (str): Output directory path for the converted TensorFlow.js model.
|
|
191
|
+
half (bool, optional): Enable FP16 quantization. Defaults to False.
|
|
192
|
+
int8 (bool, optional): Enable INT8 quantization. Defaults to False.
|
|
193
|
+
prefix (str, optional): Logging prefix. Defaults to "".
|
|
194
|
+
|
|
195
|
+
Notes:
|
|
196
|
+
Requires tensorflowjs package. Uses tensorflowjs_converter command-line tool for conversion.
|
|
197
|
+
Handles spaces in file paths and warns if output directory contains spaces.
|
|
198
|
+
"""
|
|
199
|
+
import subprocess
|
|
200
|
+
|
|
201
|
+
import tensorflow as tf
|
|
202
|
+
import tensorflowjs as tfjs
|
|
203
|
+
|
|
204
|
+
LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
|
|
205
|
+
|
|
206
|
+
gd = tf.Graph().as_graph_def() # TF GraphDef
|
|
207
|
+
with open(pb_file, "rb") as file:
|
|
208
|
+
gd.ParseFromString(file.read())
|
|
209
|
+
outputs = ",".join(gd_outputs(gd))
|
|
210
|
+
LOGGER.info(f"\n{prefix} output node names: {outputs}")
|
|
211
|
+
|
|
212
|
+
quantization = "--quantize_float16" if half else "--quantize_uint8" if int8 else ""
|
|
213
|
+
with spaces_in_path(pb_file) as fpb_, spaces_in_path(output_dir) as f_: # exporter cannot handle spaces in paths
|
|
214
|
+
cmd = (
|
|
215
|
+
"tensorflowjs_converter "
|
|
216
|
+
f'--input_format=tf_frozen_model {quantization} --output_node_names={outputs} "{fpb_}" "{f_}"'
|
|
217
|
+
)
|
|
218
|
+
LOGGER.info(f"{prefix} running '{cmd}'")
|
|
219
|
+
subprocess.run(cmd, shell=True)
|
|
220
|
+
|
|
221
|
+
if " " in output_dir:
|
|
222
|
+
LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{output_dir}'.")
|
|
223
|
+
|
|
224
|
+
|
|
225
|
+
def gd_outputs(gd):
|
|
226
|
+
"""Return TensorFlow GraphDef model output node names."""
|
|
227
|
+
name_list, input_list = [], []
|
|
228
|
+
for node in gd.node: # tensorflow.core.framework.node_def_pb2.NodeDef
|
|
229
|
+
name_list.append(node.name)
|
|
230
|
+
input_list.extend(node.input)
|
|
231
|
+
return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))
|
|
@@ -0,0 +1,219 @@
|
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import contextlib
|
|
6
|
+
import glob
|
|
7
|
+
import os
|
|
8
|
+
import shutil
|
|
9
|
+
import tempfile
|
|
10
|
+
from contextlib import contextmanager
|
|
11
|
+
from datetime import datetime
|
|
12
|
+
from pathlib import Path
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class WorkingDirectory(contextlib.ContextDecorator):
|
|
16
|
+
"""A context manager and decorator for temporarily changing the working directory.
|
|
17
|
+
|
|
18
|
+
This class allows for the temporary change of the working directory using a context manager or decorator. It ensures
|
|
19
|
+
that the original working directory is restored after the context or decorated function completes.
|
|
20
|
+
|
|
21
|
+
Attributes:
|
|
22
|
+
dir (Path | str): The new directory to switch to.
|
|
23
|
+
cwd (Path): The original current working directory before the switch.
|
|
24
|
+
|
|
25
|
+
Methods:
|
|
26
|
+
__enter__: Changes the current directory to the specified directory.
|
|
27
|
+
__exit__: Restores the original working directory on context exit.
|
|
28
|
+
|
|
29
|
+
Examples:
|
|
30
|
+
Using as a context manager:
|
|
31
|
+
>>> with WorkingDirectory("/path/to/new/dir"):
|
|
32
|
+
... # Perform operations in the new directory
|
|
33
|
+
... pass
|
|
34
|
+
|
|
35
|
+
Using as a decorator:
|
|
36
|
+
>>> @WorkingDirectory("/path/to/new/dir")
|
|
37
|
+
... def some_function():
|
|
38
|
+
... # Perform operations in the new directory
|
|
39
|
+
... pass
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
def __init__(self, new_dir: str | Path):
|
|
43
|
+
"""Initialize the WorkingDirectory context manager with the target directory."""
|
|
44
|
+
self.dir = new_dir # new dir
|
|
45
|
+
self.cwd = Path.cwd().resolve() # current dir
|
|
46
|
+
|
|
47
|
+
def __enter__(self):
|
|
48
|
+
"""Change the current working directory to the specified directory upon entering the context."""
|
|
49
|
+
os.chdir(self.dir)
|
|
50
|
+
|
|
51
|
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
52
|
+
"""Restore the original working directory when exiting the context."""
|
|
53
|
+
os.chdir(self.cwd)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@contextmanager
|
|
57
|
+
def spaces_in_path(path: str | Path):
|
|
58
|
+
"""Context manager to handle paths with spaces in their names.
|
|
59
|
+
|
|
60
|
+
If a path contains spaces, it replaces them with underscores, copies the file/directory to the new path, executes
|
|
61
|
+
the context code block, then copies the file/directory back to its original location.
|
|
62
|
+
|
|
63
|
+
Args:
|
|
64
|
+
path (str | Path): The original path that may contain spaces.
|
|
65
|
+
|
|
66
|
+
Yields:
|
|
67
|
+
(Path | str): Temporary path with any spaces replaced by underscores.
|
|
68
|
+
|
|
69
|
+
Examples:
|
|
70
|
+
>>> with spaces_in_path("/path/with spaces") as new_path:
|
|
71
|
+
... # Your code here
|
|
72
|
+
... pass
|
|
73
|
+
"""
|
|
74
|
+
# If path has spaces, replace them with underscores
|
|
75
|
+
if " " in str(path):
|
|
76
|
+
string = isinstance(path, str) # input type
|
|
77
|
+
path = Path(path)
|
|
78
|
+
|
|
79
|
+
# Create a temporary directory and construct the new path
|
|
80
|
+
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
81
|
+
tmp_path = Path(tmp_dir) / path.name.replace(" ", "_")
|
|
82
|
+
|
|
83
|
+
# Copy file/directory
|
|
84
|
+
if path.is_dir():
|
|
85
|
+
shutil.copytree(path, tmp_path)
|
|
86
|
+
elif path.is_file():
|
|
87
|
+
tmp_path.parent.mkdir(parents=True, exist_ok=True)
|
|
88
|
+
shutil.copy2(path, tmp_path)
|
|
89
|
+
|
|
90
|
+
try:
|
|
91
|
+
# Yield the temporary path
|
|
92
|
+
yield str(tmp_path) if string else tmp_path
|
|
93
|
+
|
|
94
|
+
finally:
|
|
95
|
+
# Copy file/directory back
|
|
96
|
+
if tmp_path.is_dir():
|
|
97
|
+
shutil.copytree(tmp_path, path, dirs_exist_ok=True)
|
|
98
|
+
elif tmp_path.is_file():
|
|
99
|
+
shutil.copy2(tmp_path, path) # Copy back the file
|
|
100
|
+
|
|
101
|
+
else:
|
|
102
|
+
# If there are no spaces, just yield the original path
|
|
103
|
+
yield path
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def increment_path(path: str | Path, exist_ok: bool = False, sep: str = "", mkdir: bool = False) -> Path:
|
|
107
|
+
"""Increment a file or directory path, i.e., runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
|
|
108
|
+
|
|
109
|
+
If the path exists and `exist_ok` is not True, the path will be incremented by appending a number and `sep` to the
|
|
110
|
+
end of the path. If the path is a file, the file extension will be preserved. If the path is a directory, the number
|
|
111
|
+
will be appended directly to the end of the path.
|
|
112
|
+
|
|
113
|
+
Args:
|
|
114
|
+
path (str | Path): Path to increment.
|
|
115
|
+
exist_ok (bool, optional): If True, the path will not be incremented and returned as-is.
|
|
116
|
+
sep (str, optional): Separator to use between the path and the incrementation number.
|
|
117
|
+
mkdir (bool, optional): Create a directory if it does not exist.
|
|
118
|
+
|
|
119
|
+
Returns:
|
|
120
|
+
(Path): Incremented path.
|
|
121
|
+
|
|
122
|
+
Examples:
|
|
123
|
+
Increment a directory path:
|
|
124
|
+
>>> from pathlib import Path
|
|
125
|
+
>>> path = Path("runs/exp")
|
|
126
|
+
>>> new_path = increment_path(path)
|
|
127
|
+
>>> print(new_path)
|
|
128
|
+
runs/exp2
|
|
129
|
+
|
|
130
|
+
Increment a file path:
|
|
131
|
+
>>> path = Path("runs/exp/results.txt")
|
|
132
|
+
>>> new_path = increment_path(path)
|
|
133
|
+
>>> print(new_path)
|
|
134
|
+
runs/exp/results2.txt
|
|
135
|
+
"""
|
|
136
|
+
path = Path(path) # os-agnostic
|
|
137
|
+
if path.exists() and not exist_ok:
|
|
138
|
+
path, suffix = (path.with_suffix(""), path.suffix) if path.is_file() else (path, "")
|
|
139
|
+
|
|
140
|
+
# Method 1
|
|
141
|
+
for n in range(2, 9999):
|
|
142
|
+
p = f"{path}{sep}{n}{suffix}" # increment path
|
|
143
|
+
if not os.path.exists(p):
|
|
144
|
+
break
|
|
145
|
+
path = Path(p)
|
|
146
|
+
|
|
147
|
+
if mkdir:
|
|
148
|
+
path.mkdir(parents=True, exist_ok=True) # make directory
|
|
149
|
+
|
|
150
|
+
return path
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
def file_age(path: str | Path = __file__) -> int:
|
|
154
|
+
"""Return days since the last modification of the specified file."""
|
|
155
|
+
dt = datetime.now() - datetime.fromtimestamp(Path(path).stat().st_mtime) # delta
|
|
156
|
+
return dt.days # + dt.seconds / 86400 # fractional days
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def file_date(path: str | Path = __file__) -> str:
|
|
160
|
+
"""Return the file modification date in 'YYYY-M-D' format."""
|
|
161
|
+
t = datetime.fromtimestamp(Path(path).stat().st_mtime)
|
|
162
|
+
return f"{t.year}-{t.month}-{t.day}"
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
def file_size(path: str | Path) -> float:
|
|
166
|
+
"""Return the size of a file or directory in megabytes (MB)."""
|
|
167
|
+
if isinstance(path, (str, Path)):
|
|
168
|
+
mb = 1 << 20 # bytes to MiB (1024 ** 2)
|
|
169
|
+
path = Path(path)
|
|
170
|
+
if path.is_file():
|
|
171
|
+
return path.stat().st_size / mb
|
|
172
|
+
elif path.is_dir():
|
|
173
|
+
return sum(f.stat().st_size for f in path.glob("**/*") if f.is_file()) / mb
|
|
174
|
+
return 0.0
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
def get_latest_run(search_dir: str = ".") -> str:
|
|
178
|
+
"""Return the path to the most recent 'last.pt' file in the specified directory for resuming training."""
|
|
179
|
+
last_list = glob.glob(f"{search_dir}/**/last*.pt", recursive=True)
|
|
180
|
+
return max(last_list, key=os.path.getctime) if last_list else ""
|
|
181
|
+
|
|
182
|
+
|
|
183
|
+
def update_models(model_names: tuple = ("yolo11n.pt",), source_dir: Path = Path("."), update_names: bool = False):
|
|
184
|
+
"""Update and re-save specified YOLO models in an 'updated_models' subdirectory.
|
|
185
|
+
|
|
186
|
+
Args:
|
|
187
|
+
model_names (tuple, optional): Model filenames to update.
|
|
188
|
+
source_dir (Path, optional): Directory containing models and target subdirectory.
|
|
189
|
+
update_names (bool, optional): Update model names from a data YAML.
|
|
190
|
+
|
|
191
|
+
Examples:
|
|
192
|
+
Update specified YOLO models and save them in 'updated_models' subdirectory:
|
|
193
|
+
>>> from ultralytics.utils.files import update_models
|
|
194
|
+
>>> model_names = ("yolo11n.pt", "yolov8s.pt")
|
|
195
|
+
>>> update_models(model_names, source_dir=Path("/models"), update_names=True)
|
|
196
|
+
"""
|
|
197
|
+
from ultralytics import YOLO
|
|
198
|
+
from ultralytics.nn.autobackend import default_class_names
|
|
199
|
+
from ultralytics.utils import LOGGER
|
|
200
|
+
|
|
201
|
+
target_dir = source_dir / "updated_models"
|
|
202
|
+
target_dir.mkdir(parents=True, exist_ok=True) # Ensure target directory exists
|
|
203
|
+
|
|
204
|
+
for model_name in model_names:
|
|
205
|
+
model_path = source_dir / model_name
|
|
206
|
+
LOGGER.info(f"Loading model from {model_path}")
|
|
207
|
+
|
|
208
|
+
# Load model
|
|
209
|
+
model = YOLO(model_path)
|
|
210
|
+
model.half()
|
|
211
|
+
if update_names: # update model names from a dataset YAML
|
|
212
|
+
model.model.names = default_class_names("coco8.yaml")
|
|
213
|
+
|
|
214
|
+
# Define new save path
|
|
215
|
+
save_path = target_dir / model_name
|
|
216
|
+
|
|
217
|
+
# Save model using model.save()
|
|
218
|
+
LOGGER.info(f"Re-saving {model_name} model to {save_path}")
|
|
219
|
+
model.save(save_path)
|
ultralytics/utils/git.py
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from functools import cached_property
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class GitRepo:
|
|
10
|
+
"""Represent a local Git repository and expose branch, commit, and remote metadata.
|
|
11
|
+
|
|
12
|
+
This class discovers the repository root by searching for a .git entry from the given path upward, resolves the
|
|
13
|
+
actual .git directory (including worktrees), and reads Git metadata directly from on-disk files. It does not invoke
|
|
14
|
+
the git binary and therefore works in restricted environments. All metadata properties are resolved lazily and
|
|
15
|
+
cached; construct a new instance to refresh state.
|
|
16
|
+
|
|
17
|
+
Attributes:
|
|
18
|
+
root (Path | None): Repository root directory containing the .git entry; None if not in a repository.
|
|
19
|
+
gitdir (Path | None): Resolved .git directory path; handles worktrees; None if unresolved.
|
|
20
|
+
head (str | None): Raw contents of HEAD; a SHA for detached HEAD or "ref: <refname>" for branch heads.
|
|
21
|
+
is_repo (bool): Whether the provided path resides inside a Git repository.
|
|
22
|
+
branch (str | None): Current branch name when HEAD points to a branch; None for detached HEAD or non-repo.
|
|
23
|
+
commit (str | None): Current commit SHA for HEAD; None if not determinable.
|
|
24
|
+
origin (str | None): URL of the "origin" remote as read from gitdir/config; None if unset or unavailable.
|
|
25
|
+
|
|
26
|
+
Examples:
|
|
27
|
+
Initialize from the current working directory and read metadata
|
|
28
|
+
>>> from pathlib import Path
|
|
29
|
+
>>> repo = GitRepo(Path.cwd())
|
|
30
|
+
>>> repo.is_repo
|
|
31
|
+
True
|
|
32
|
+
>>> repo.branch, repo.commit[:7], repo.origin
|
|
33
|
+
('main', '1a2b3c4', 'https://example.com/owner/repo.git')
|
|
34
|
+
|
|
35
|
+
Notes:
|
|
36
|
+
- Resolves metadata by reading files: HEAD, packed-refs, and config; no subprocess calls are used.
|
|
37
|
+
- Caches properties on first access using cached_property; recreate the object to reflect repository changes.
|
|
38
|
+
"""
|
|
39
|
+
|
|
40
|
+
def __init__(self, path: Path = Path(__file__).resolve()):
|
|
41
|
+
"""Initialize a Git repository context by discovering the repository root from a starting path.
|
|
42
|
+
|
|
43
|
+
Args:
|
|
44
|
+
path (Path, optional): File or directory path used as the starting point to locate the repository root.
|
|
45
|
+
"""
|
|
46
|
+
self.root = self._find_root(path)
|
|
47
|
+
self.gitdir = self._gitdir(self.root) if self.root else None
|
|
48
|
+
|
|
49
|
+
@staticmethod
|
|
50
|
+
def _find_root(p: Path) -> Path | None:
|
|
51
|
+
"""Return repo root or None."""
|
|
52
|
+
return next((d for d in [p, *list(p.parents)] if (d / ".git").exists()), None)
|
|
53
|
+
|
|
54
|
+
@staticmethod
|
|
55
|
+
def _gitdir(root: Path) -> Path | None:
|
|
56
|
+
"""Resolve actual .git directory (handles worktrees)."""
|
|
57
|
+
g = root / ".git"
|
|
58
|
+
if g.is_dir():
|
|
59
|
+
return g
|
|
60
|
+
if g.is_file():
|
|
61
|
+
t = g.read_text(errors="ignore").strip()
|
|
62
|
+
if t.startswith("gitdir:"):
|
|
63
|
+
return (root / t.split(":", 1)[1].strip()).resolve()
|
|
64
|
+
return None
|
|
65
|
+
|
|
66
|
+
@staticmethod
|
|
67
|
+
def _read(p: Path | None) -> str | None:
|
|
68
|
+
"""Read and strip file if exists."""
|
|
69
|
+
return p.read_text(errors="ignore").strip() if p and p.exists() else None
|
|
70
|
+
|
|
71
|
+
@cached_property
|
|
72
|
+
def head(self) -> str | None:
|
|
73
|
+
"""HEAD file contents."""
|
|
74
|
+
return self._read(self.gitdir / "HEAD" if self.gitdir else None)
|
|
75
|
+
|
|
76
|
+
def _ref_commit(self, ref: str) -> str | None:
|
|
77
|
+
"""Commit for ref (handles packed-refs)."""
|
|
78
|
+
rf = self.gitdir / ref
|
|
79
|
+
if s := self._read(rf):
|
|
80
|
+
return s
|
|
81
|
+
pf = self.gitdir / "packed-refs"
|
|
82
|
+
b = pf.read_bytes().splitlines() if pf.exists() else []
|
|
83
|
+
tgt = ref.encode()
|
|
84
|
+
for line in b:
|
|
85
|
+
if line[:1] in (b"#", b"^") or b" " not in line:
|
|
86
|
+
continue
|
|
87
|
+
sha, name = line.split(b" ", 1)
|
|
88
|
+
if name.strip() == tgt:
|
|
89
|
+
return sha.decode()
|
|
90
|
+
return None
|
|
91
|
+
|
|
92
|
+
@property
|
|
93
|
+
def is_repo(self) -> bool:
|
|
94
|
+
"""True if inside a git repo."""
|
|
95
|
+
return self.gitdir is not None
|
|
96
|
+
|
|
97
|
+
@cached_property
|
|
98
|
+
def branch(self) -> str | None:
|
|
99
|
+
"""Current branch or None."""
|
|
100
|
+
if not self.is_repo or not self.head or not self.head.startswith("ref: "):
|
|
101
|
+
return None
|
|
102
|
+
ref = self.head[5:].strip()
|
|
103
|
+
return ref[len("refs/heads/") :] if ref.startswith("refs/heads/") else ref
|
|
104
|
+
|
|
105
|
+
@cached_property
|
|
106
|
+
def commit(self) -> str | None:
|
|
107
|
+
"""Current commit SHA or None."""
|
|
108
|
+
if not self.is_repo or not self.head:
|
|
109
|
+
return None
|
|
110
|
+
return self._ref_commit(self.head[5:].strip()) if self.head.startswith("ref: ") else self.head
|
|
111
|
+
|
|
112
|
+
@cached_property
|
|
113
|
+
def origin(self) -> str | None:
|
|
114
|
+
"""Origin URL or None."""
|
|
115
|
+
if not self.is_repo:
|
|
116
|
+
return None
|
|
117
|
+
cfg = self.gitdir / "config"
|
|
118
|
+
remote, url = None, None
|
|
119
|
+
for s in (self._read(cfg) or "").splitlines():
|
|
120
|
+
t = s.strip()
|
|
121
|
+
if t.startswith("[") and t.endswith("]"):
|
|
122
|
+
remote = t.lower()
|
|
123
|
+
elif t.lower().startswith("url =") and remote == '[remote "origin"]':
|
|
124
|
+
url = t.split("=", 1)[1].strip()
|
|
125
|
+
break
|
|
126
|
+
return url
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
if __name__ == "__main__":
|
|
130
|
+
import time
|
|
131
|
+
|
|
132
|
+
g = GitRepo()
|
|
133
|
+
if g.is_repo:
|
|
134
|
+
t0 = time.perf_counter()
|
|
135
|
+
print(f"repo={g.root}\nbranch={g.branch}\ncommit={g.commit}\norigin={g.origin}")
|
|
136
|
+
dt = (time.perf_counter() - t0) * 1000
|
|
137
|
+
print(f"\n⏱️ Profiling: total {dt:.3f} ms")
|