83 lines
2.2 KiB
Python
83 lines
2.2 KiB
Python
import cv2
|
|
|
|
from numpy import ones, uint8, concatenate
|
|
|
|
from ...robot.controller import SensorsController
|
|
from ...robot.sensor import Sensor
|
|
|
|
|
|
class Blob(Sensor):
|
|
registers = Sensor.registers + ['center', 'radius']
|
|
|
|
def __init__(self, x, y, radius):
|
|
self.center = x, y
|
|
self.radius = radius
|
|
|
|
def draw(self, img, color=(255, 0, 0), thickness=3):
|
|
cv2.circle(img, self.center, self.radius, color, thickness)
|
|
|
|
@property
|
|
def json(self):
|
|
return {"center": self.center, "radius": self.radius}
|
|
|
|
|
|
class BlobDetector(SensorsController):
|
|
channels = {
|
|
'R': 2, 'G': 1, 'B': 0,
|
|
'H': 0, 'S': 1, 'V': 2
|
|
}
|
|
|
|
def __init__(self, robot, name, cameras, freq, filters):
|
|
SensorsController.__init__(self, None, [], freq)
|
|
|
|
self.name = name
|
|
|
|
self._robot = robot
|
|
self._names = cameras
|
|
self._blobs = []
|
|
self.filters = filters
|
|
|
|
def detect_blob(self, img, filters):
|
|
"""
|
|
"filters" must be something similar to:
|
|
filters = {
|
|
'R': (150, 255), # (min, max)
|
|
'S': (150, 255),
|
|
}
|
|
|
|
"""
|
|
acc_mask = ones(img.shape[:2], dtype=uint8) * 255
|
|
|
|
rgb = img.copy()
|
|
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
|
|
|
for c, (min, max) in filters.items():
|
|
img = rgb if c in 'RGB' else hsv
|
|
|
|
mask = img[:, :, self.channels[c]]
|
|
mask[mask < min] = 0
|
|
mask[mask > max] = 0
|
|
|
|
acc_mask &= mask
|
|
|
|
kernel = ones((5, 5), uint8)
|
|
acc_mask = cv2.dilate(cv2.erode(acc_mask, kernel), kernel)
|
|
|
|
circles = cv2.HoughCircles(acc_mask, cv2.HOUGH_GRADIENT, 3, img.shape[0] / 5.)
|
|
return circles.reshape(-1, 3) if circles is not None else []
|
|
|
|
def update(self):
|
|
if not hasattr(self, 'cameras'):
|
|
self.cameras = [getattr(self._robot, c) for c in self._names]
|
|
|
|
self._blobs = concatenate([self.detect_blob(c.frame, self.filters)
|
|
for c in self.cameras])
|
|
|
|
@property
|
|
def blobs(self):
|
|
return [Blob(*b) for b in self._blobs]
|
|
|
|
@property
|
|
def registers(self):
|
|
return ['blobs']
|