Code: Select all
# python
import numpy as np
import cv2 as cv
def readFrame():
sim.handleVisionSensor(self.handle)
buf, res = sim.getVisionSensorImg(self.handle)
return np.frombuffer(buf, dtype=np.uint8).reshape(*res, 3)
def showImage(k, img):
sim.setVisionSensorImg(sim.getObject(f':/{k}'), img.tobytes())
def sysCall_init():
sim = require('sim')
self.handle = sim.getObject('.')
robot = sim.getObject(':')
script = sim.getScript(sim.scripttype_childscript, robot)
self.robotScriptFuncs = sim.getScriptFunctions(script)
def sysCall_thread():
while True:
frame = readFrame()
imgw, imgh, imgc = frame.shape
blurred = cv.GaussianBlur(frame, (13, 13), 0)
hsv = cv.cvtColor(blurred, cv.COLOR_RGB2HSV)
# construct a mask for the color "red", then perform
# a series of dilations and erosions to remove any small
# blobs left in the mask
mask = cv.inRange(hsv, (0, 200, 64), (80, 255, 200))
n = 5
mask = cv.dilate(mask, None, iterations=n)
mask = cv.erode(mask, None, iterations=n)
maskdbg = cv.cvtColor(mask, cv.COLOR_GRAY2RGB)
contours, hierarchy = cv.findContours(mask.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
if contours:
c = max(contours, key = cv.contourArea)
x, y, w, h = cv.boundingRect(c)
cv.rectangle(maskdbg, (x, y), (x+w, y+h), (0, 255, 0), 2)
cv.arrowedLine(maskdbg, (x+w//2, y+h//2), (imgw//2, imgh//2), (255, 0, 0), 2)
# normalize rect in -1 ... +1:
normalize = lambda x: np.interp(x[0], (0, x[1]), (-1, 1))
x, y, w, h = map(normalize, ((x, imgw), (y, imgh), (w, imgw), (h, imgh)))
# box center:
self.robotScriptFuncs.setBoxCenter(x + w/2, y + h/2)
showImage('blurred', blurred)
showImage('hsv', hsv)
showImage('mask', maskdbg)
Thank you in advance!!!