157 lines
5.1 KiB
Python
157 lines
5.1 KiB
Python
"""
|
|
Beer Bot Vision - Hauptklasse für die Bildverarbeitung mit YOLOv8.
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from config import Config
|
|
from calibrator import CameraCalibrator
|
|
from user_interface import ModernUI
|
|
|
|
# YOLOv8 Import
|
|
try:
|
|
from ultralytics import YOLO
|
|
YOLO_AVAILABLE = True
|
|
except ImportError:
|
|
YOLO_AVAILABLE = False
|
|
print("⚠️ ultralytics nicht installiert. Installiere mit: pip install ultralytics")
|
|
|
|
|
|
class BeerBotVision:
|
|
def __init__(self, camera_id=None, model_path=None):
|
|
self.camera_id = Config.CAMERA_ID if camera_id is None else camera_id
|
|
self.model_path = model_path or Config.YOLO_MODEL
|
|
|
|
if YOLO_AVAILABLE:
|
|
print(f"🔄 Lade YOLO-Modell: {self.model_path}")
|
|
self.model = YOLO(self.model_path)
|
|
print("✅ YOLO-Modell geladen!")
|
|
else:
|
|
self.model = None
|
|
|
|
self.cap = None
|
|
self.camera_initialized = False
|
|
self.calibrator = CameraCalibrator()
|
|
self.detected_bottles = []
|
|
self.ui = ModernUI()
|
|
|
|
def init_camera(self):
|
|
print(f"🔄 Öffne Kamera {self.camera_id}...")
|
|
self.cap = cv2.VideoCapture(self.camera_id)
|
|
if not self.cap.isOpened():
|
|
print("❌ Kamera konnte nicht geöffnet werden!")
|
|
return False
|
|
self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, Config.CAMERA_WIDTH)
|
|
self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, Config.CAMERA_HEIGHT)
|
|
self.camera_initialized = True
|
|
print(f"✅ Kamera initialisiert")
|
|
return True
|
|
|
|
def capture_frame(self):
|
|
if not self.camera_initialized:
|
|
if not self.init_camera():
|
|
return None
|
|
ret, frame = self.cap.read()
|
|
return frame if ret else None
|
|
|
|
def detect_objects(self, frame=None, visualize=False):
|
|
if self.model is None:
|
|
return {"bottles": []}
|
|
|
|
if frame is None:
|
|
frame = self.capture_frame()
|
|
if frame is None:
|
|
return {"bottles": []}
|
|
|
|
results = self.model(frame, verbose=False)
|
|
self.detected_bottles = []
|
|
|
|
for result in results:
|
|
for box in result.boxes:
|
|
cls = int(box.cls[0])
|
|
conf = float(box.conf[0])
|
|
|
|
if conf >= Config.CONFIDENCE_THRESHOLD and cls == Config.BOTTLE_CLASS_ID:
|
|
x1, y1, x2, y2 = map(int, box.xyxy[0])
|
|
center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
|
|
|
|
robot_coords = None
|
|
if self.calibrator.is_calibrated:
|
|
try:
|
|
robot_coords = self.calibrator.pixel_to_robot(center_x, center_y)
|
|
except:
|
|
pass
|
|
|
|
self.detected_bottles.append({
|
|
"bbox": (x1, y1, x2, y2),
|
|
"center_pixel": (center_x, center_y),
|
|
"robot_coords": robot_coords,
|
|
"confidence": conf
|
|
})
|
|
|
|
if visualize:
|
|
self._visualize_detections(frame)
|
|
|
|
return {"bottles": self.detected_bottles, "frame": frame}
|
|
|
|
def _visualize_detections(self, frame):
|
|
for bottle in self.detected_bottles:
|
|
x1, y1, x2, y2 = bottle["bbox"]
|
|
cx, cy = bottle["center_pixel"]
|
|
conf = bottle["confidence"]
|
|
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
cv2.circle(frame, (cx, cy), 5, (0, 0, 255), -1)
|
|
|
|
label = f"Bottle {conf:.0%}"
|
|
if bottle["robot_coords"]:
|
|
rx, ry = bottle["robot_coords"]
|
|
label += f" ({rx:.0f}, {ry:.0f})"
|
|
|
|
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
|
|
|
h, w = frame.shape[:2]
|
|
cv2.rectangle(frame, (0, h - 35), (w, h), (0, 0, 0), -1)
|
|
cv2.putText(frame, f"Flaschen: {len(self.detected_bottles)} | Druecke 0 fuer Hauptmenue",
|
|
(10, h - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
|
|
|
|
cv2.imshow("Beer Bot Vision", frame)
|
|
|
|
def get_nearest_bottle(self, reference_point=(0, 0)):
|
|
if not self.detected_bottles:
|
|
self.detect_objects()
|
|
|
|
valid = [b for b in self.detected_bottles if b["robot_coords"]]
|
|
if not valid:
|
|
return None
|
|
|
|
def dist(b):
|
|
bx, by = b["robot_coords"]
|
|
rx, ry = reference_point
|
|
return np.sqrt((bx - rx) ** 2 + (by - ry) ** 2)
|
|
|
|
return min(valid, key=dist)
|
|
|
|
def run_live_detection(self):
|
|
print("🎥 Live-Erkennung gestartet. Drücke '0' zum Beenden.")
|
|
if not self.camera_initialized:
|
|
self.init_camera()
|
|
|
|
while True:
|
|
frame = self.capture_frame()
|
|
if frame is None:
|
|
break
|
|
self.detect_objects(frame, visualize=True)
|
|
if cv2.waitKey(1) & 0xFF == ord('0'):
|
|
break
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
def load_calibration(self, matrix_file=None):
|
|
return self.calibrator.load_calibration(matrix_file)
|
|
|
|
def release(self):
|
|
if self.cap:
|
|
self.cap.release()
|
|
cv2.destroyAllWindows()
|