Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
"""Distance estimation."""
import PIL
import SimpleCV
import io
import math
import picamera
import picamera.array
import signal
import sys
import time
class PictureWrapper(object):
def __init__(self, resolution=None):
self._camera = picamera.PiCamera()
scale = 5
if resolution is not None:
self._camera.resolution = resolution
else:
self._camera.resolution = (320, 240)
self._camera.start_preview()
def get_image(self):
try:
with picamera.array.PiRGBArray(self._camera) as stream:
self._camera.capture(stream, format='rgb', use_video_port=True)
pil = PIL.Image.fromarray(stream.array, mode='RGB')
image = SimpleCV.Image(pil)
return image
# This usually happens when we try to shut down the program using
# SIGINT; just ignore it
except (picamera.PiCameraValueError, picamera.PiCameraAlreadyRecording):
return None
# pylint: disable=superfluous-parens
# pylint: disable=global-statement
class Watch(object):
""""Does image processing to determine the distance and direction to a
target.
"""
HORIZONTAL_VIEWPORT_D = 55 * 2.0
HORIZONTAL_VIEWPORT_R = math.radians(HORIZONTAL_VIEWPORT_D)
ORANGE_DUCT_HUE = 7.0
CD_DIAMETER_M = .110
def __init__(self, picture_wrapper):
self._wrapper = picture_wrapper
image = picture_wrapper.get_image()
self._r_per_pixel = self.HORIZONTAL_VIEWPORT_R / image.width
self._d_per_pixel = self.HORIZONTAL_VIEWPORT_D / image.width
# From manual testing, 100 seems to be a good minsize for 320x240
self._min_size = int(image.width * image.height / 320.0 / 240.0 * 100.0)
def _get_distance_m_direction_d(self, image):
"""Returns the distance in meters and the direction in degrees to the
target from an image.
"""
if image is None:
return (None, None)
crop_y = int(image.height * .3)
crop_height = int(image.height * .5)
def crop(img):
return img.crop(0, crop_y, image.width, crop_height)
# The hue distance to orange from black is pretty low, so we need
# to create a black mask and subtract those parts first
# TODO: I don't know why, but cropping the image first and then doing
# distance produces an image with weird streaks instead of anything
# correct, so we can't just crop it immediately
#image = crop(image)
#black_mask = image.colorDistance(SimpleCV.Color.BLACK).binarize()
#distance = image.hueDistance(self.ORANGE_DUCT_HUE).invert().binarize(220).invert()
black_mask = crop(image.colorDistance(SimpleCV.Color.BLACK)).binarize()
distance = crop(image).hueDistance(self.ORANGE_DUCT_HUE).invert().binarize(220).invert()
blobs = (distance - black_mask).findBlobs(
minsize=self._min_size,
maxsize=(distance.width * distance.height // 4)
)
if blobs is not None:
circles = [b for b in blobs if b.isCircle(0.5)]
if len(circles) == 0:
return (None, None)
# The circles should already be sorted by area ascending
angle_r = circles[-1].radius() * self._r_per_pixel
distance_m = self.CD_DIAMETER_M / math.tan(angle_r)
mid_circle_x = circles[-1].x + circles[-1].radius() * 0.5
angle_d = (mid_circle_x - distance.width * 0.5) * self._d_per_pixel
if angle_d < 0.0:
angle_d += 360.0
return (distance_m, angle_d)
return (None, None)
def get_distance_m_direction_d(self):
"""Returns the distance in meters and the direction in degrees to the
target.
"""
image = self._wrapper.get_image()
return self._get_distance_m_direction_d(image)
START = None
ITERATIONS = None
END = None
def terminate(signal_number, stack_frame):
print(
'{iterations} iterations took {seconds} seconds, fps = {fps}'.format(
iterations=ITERATIONS,
seconds=(END - START),
fps=(ITERATIONS / (END - START))
)
)
sys.exit()
def main():
signal.signal(signal.SIGINT, terminate)
picture_wrapper = PictureWrapper()
watch = Watch(picture_wrapper)
global START
global ITERATIONS
global END
ITERATIONS = 0
START = time.time()
while True:
distance_m, angle_d = watch.get_distance_m_direction_d()
print(
'Distance: {distance_m} m {distance_in} in, angle: {angle}'.format(
distance_m=distance_m,
distance_in=(None if distance_m is None else distance_m * 39.3700787),
angle=angle_d
)
)
ITERATIONS += 1
END = time.time()
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment