Exemplo n.º 1
0
class CupDetector(event.DecisionMaker):
    def __init__(self, ev, cam_angle, cup_color='pahar_mare_albastru'):
        self.frames_seen = 0
        self.cam_angle = cam_angle
        self.cup_color = cup_color
        self.blue_cup = ColorMatcher(cup_color)
        self.ev = ev
        super(CupDetector, self).__init__(ev)

    def frame(self, frame):
        big_contours = self.blue_cup.find_bboxes(frame)

        contours = []
        for contour in big_contours:
            x, y, X, Y = contour
            ratio = float(Y-y)/(X-x+1)
            contours.append((x, y, X, Y, 1, 1.2))

        for x, y, X, Y in big_contours:
            ratio = float(Y-y)/(X-x+1)
            cv2.rectangle(frame, (x-2, y-2), (X, Y), (255, 0, 0), 2)
            cv2.putText(frame, '%0.3f' % ratio, (x, y+20),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255),
                        thickness=2)
        coords_list = []
        for x, y, X, Y, matches, ratio in contours:
            cv2.rectangle(frame, (x - 2, y - 2), (X, Y), (0, 255, 0), 2)
            dist = '%0.2f' % get_distance_from_cup_width(X-x)
            coords = pixels2coords((x+X)/2., Y-(X-x), X-x,
                                   cam_angle=self.cam_angle)
            cv2.putText(frame, dist, (x, y-20),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255),
                        thickness=2)
            cv2.putText(frame, '%0.2f %0.2f %0.2f' % coords, (x, y-50),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255),
                        thickness=2)
            if x > 0 and X < frame.shape[1]:
                coords_list.append(coords)
            coords_list.sort()
            for x, y, z in coords_list:
                self.frames_seen = min(self.frames_seen + 1, 20)
                if self.frames_seen == 20 and x < 400:
                    print 'cd: Cup appeared: %s' % self.cup_color
                    self.emit('cup_appeared', (x, y, z))
                    self.frames_seen = 0
                    break
        #else:
            #print 'cd: Cups done: %s' % self.cup_color
            #self.emit('cups_done')

        cv2.imshow('Cup detector', cv2.resize(frame, dsize=None,
                                              fx=0.5, fy=0.5))