62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
from ultralytics import YOLO
|
|
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("source", help="Фото, видео, папка с фото или 0 для веб-камеры")
|
|
parser.add_argument("--weights", default=str(HERE / "models/latest.pt"))
|
|
parser.add_argument("--conf", type=float, default=0.25)
|
|
args = parser.parse_args()
|
|
|
|
source = int(args.source) if args.source.isdigit() else args.source
|
|
model = YOLO(args.weights)
|
|
log_path = HERE / "logs/simple_module_b.log"
|
|
log_path.parent.mkdir(exist_ok=True)
|
|
|
|
still_images = False
|
|
if isinstance(source, str):
|
|
path = Path(source)
|
|
still_images = path.is_dir() or path.suffix.lower() in {
|
|
".jpg",
|
|
".jpeg",
|
|
".png",
|
|
".bmp",
|
|
".webp",
|
|
}
|
|
|
|
with log_path.open("a", encoding="utf-8") as log:
|
|
results = model.predict(source=source, conf=args.conf, stream=True, verbose=False)
|
|
for result in results:
|
|
now = datetime.now().isoformat(timespec="seconds")
|
|
detections = []
|
|
if result.boxes is not None:
|
|
for box in result.boxes:
|
|
name = result.names[int(box.cls.item())]
|
|
confidence = float(box.conf.item())
|
|
detections.append(f"{name} {confidence:.2f}")
|
|
|
|
line = f"{now} | {', '.join(detections) if detections else 'ничего не найдено'}"
|
|
print(line)
|
|
log.write(line + "\n")
|
|
log.flush()
|
|
|
|
cv2.imshow("YOLO: q или Esc для выхода", result.plot())
|
|
key = cv2.waitKey(0 if still_images else 1) & 0xFF
|
|
if key in (ord("q"), 27):
|
|
break
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|