adding real-time contact area extraction example

Author: vocdex

examples/contact_area.py

@@ -29,5 +29,33 @@ def extract_surface_contact():
     ImageHandler.save("surface_contact_2.png", sample_img_2)
 
 
+def extract_surface_contact_real_time():
+    """
+    This function demonstrates how to use the ContactArea task with real-time mode with DIGIT sensor.
+    Make sure to have a DIGIT sensor connected to your computer and setup the digit-interface library.
+    Output:
+        Displays a real-time object contact area by fitting an ellipse.
+    """
+    from digit_interface import Digit
+    import cv2
+    base_img_path = "/home/shuk/digits2/images/0000/0041.png"
+    base_img = ImageHandler(base_img_path).nparray
+    # setting up the DIGIT sensor
+    digit = Digit("Dxxxxx")  # replace xxxxx with your DIGIT serial number
+    digit.connect()
+    while True:
+        frame=digit.get_frame()
+        img=cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # target image frame for __call__ method
+        contact_area=ContactArea(base=base_img, contour_threshold=10,real_time=True)
+        contact_area(target=img) # __call__ method
+        cv2.imshow('surface contact', img)
+        k=cv2.waitKey(1)
+        if k%256==27:
+            #ESC hit
+            print("Escape hit, closing...")
+            break
+
+
 if __name__ == "__main__":
-    extract_surface_contact()
+    #extract_surface_contact()
+    extract_surface_contact_real_time()

pytouch/tasks/contact_area.py

@@ -10,11 +10,12 @@
 
 class ContactArea:
     def __init__(
-        self, base=None, draw_poly=True, contour_threshold=100, *args, **kwargs
+        self, base=None, draw_poly=True, contour_threshold=100,real_time=False,*args, **kwargs
     ):
         self.base = base
         self.draw_poly = draw_poly
         self.contour_threshold = contour_threshold
+        self.real_time = real_time
 
     def __call__(self, target, base=None):
         base = self.base if base is None else base
@@ -23,13 +24,18 @@ def __call__(self, target, base=None):
         diff = self._diff(target, base)
         diff = self._smooth(diff)
         contours = self._contours(diff)
-        (
-            poly,
-            major_axis,
-            major_axis_end,
-            minor_axis,
-            minor_axis_end,
-        ) = self._compute_contact_area(contours, self.contour_threshold)
+        if self._compute_contact_area(contours, self.contour_threshold) == None and self.real_time==False:
+            raise Exception("No contact area detected.")
+        if self._compute_contact_area(contours, self.contour_threshold) == None and self.real_time==True:
+            return None
+        else:
+            (
+                poly,
+                major_axis,
+                major_axis_end,
+                minor_axis,
+                minor_axis_end,
+            ) = self._compute_contact_area(contours, self.contour_threshold)
         if self.draw_poly:
             self._draw_major_minor(
                 target, poly, major_axis, major_axis_end, minor_axis, minor_axis_end
@@ -104,4 +110,4 @@ def _compute_contact_area(self, contours, contour_threshold):
                 )
                 major_axis_end = 2 * center - major_axis
                 minor_axis_end = 2 * center - minor_axis
-        return poly, major_axis, major_axis_end, minor_axis, minor_axis_end
+                return poly, major_axis, major_axis_end, minor_axis, minor_axis_end