SHA256
150 lines
3.3 KiB
Python
150 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
|
IMAGE_PATH = Path("/dev/shm/robot_frame.jpg")
|
|
OUTPUT_PATH = Path("/dev/shm/r1_detections.json")
|
|
|
|
MODEL = "yolo26n.pt"
|
|
|
|
CONFIDENCE = 0.35
|
|
IMAGE_SIZE = 640
|
|
|
|
|
|
def atomic_write_json(path: Path, data):
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w",
|
|
dir=path.parent,
|
|
delete=False,
|
|
suffix=".tmp",
|
|
) as f:
|
|
json.dump(
|
|
data,
|
|
f,
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
|
|
tmp = f.name
|
|
|
|
os.replace(tmp, path)
|
|
|
|
|
|
def main():
|
|
print(f"[YOLO] Loading {MODEL}")
|
|
|
|
model = YOLO(MODEL)
|
|
|
|
print("[YOLO] Ready")
|
|
print(f"[YOLO] Camera: {IMAGE_PATH}")
|
|
print(f"[YOLO] Output: {OUTPUT_PATH}")
|
|
|
|
last_mtime = None
|
|
|
|
while True:
|
|
try:
|
|
if not IMAGE_PATH.exists():
|
|
time.sleep(0.05)
|
|
continue
|
|
|
|
mtime = IMAGE_PATH.stat().st_mtime_ns
|
|
|
|
if mtime == last_mtime:
|
|
time.sleep(0.01)
|
|
continue
|
|
|
|
last_mtime = mtime
|
|
|
|
results = model.predict(
|
|
source=str(IMAGE_PATH),
|
|
conf=CONFIDENCE,
|
|
imgsz=IMAGE_SIZE,
|
|
verbose=False,
|
|
)
|
|
|
|
if not results:
|
|
continue
|
|
|
|
result = results[0]
|
|
|
|
height, width = result.orig_shape
|
|
|
|
detections = []
|
|
|
|
if result.boxes is not None:
|
|
for box in result.boxes:
|
|
class_id = int(box.cls[0])
|
|
confidence = float(box.conf[0])
|
|
|
|
x1, y1, x2, y2 = [
|
|
float(v)
|
|
for v in box.xyxy[0]
|
|
]
|
|
|
|
cx = (x1 + x2) / 2.0
|
|
cy = (y1 + y2) / 2.0
|
|
|
|
detections.append({
|
|
"class_id": class_id,
|
|
"class": result.names[class_id],
|
|
"confidence": confidence,
|
|
|
|
"box": {
|
|
"x1": x1,
|
|
"y1": y1,
|
|
"x2": x2,
|
|
"y2": y2,
|
|
},
|
|
|
|
"center": {
|
|
"x": cx,
|
|
"y": cy,
|
|
},
|
|
|
|
"normalized_center": {
|
|
"x": cx / width,
|
|
"y": cy / height,
|
|
},
|
|
})
|
|
|
|
data = {
|
|
"timestamp": time.time(),
|
|
"width": width,
|
|
"height": height,
|
|
"detections": detections,
|
|
}
|
|
|
|
atomic_write_json(
|
|
OUTPUT_PATH,
|
|
data,
|
|
)
|
|
|
|
objects = ", ".join(
|
|
f"{d['class']} {d['confidence']:.2f}"
|
|
for d in detections
|
|
)
|
|
|
|
if objects:
|
|
print(f"[YOLO] {objects}")
|
|
else:
|
|
print("[YOLO] no objects")
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n[YOLO] stopped")
|
|
break
|
|
|
|
except Exception as e:
|
|
print(f"[YOLO] error: {e}")
|
|
time.sleep(0.25)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|