Fixes need to work with ARDrone 2.0 including deal

.gitignore

@@ -0,0 +1,3 @@
+*.pyc
+*~
+.cache

.travis.yml

@@ -0,0 +1,11 @@
+language: python
+python:
+  - "2.7"
+notifications:
+  - email: false
+install:
+  - sudo apt-get update
+  - sudo apt-get install ffmpeg
+  - pip install -r requirements.txt
+  - pip install -r dev-requirements.txt
+script: py.test -n 4

README

@@ -43,9 +43,7 @@ Requirements:
 This software was tested with the following setup:
 
   * Python 2.6.6
-  * Psyco 1.6 (recommended)
-  * Pygame 1.8.1 (only for the demo)
-  * Unmodified AR.Drone firmware 1.5.1
+  * Unmodified AR.Drone firmware 2.0
 
 
 License:

dev-requirements.txt

@@ -0,0 +1,2 @@
+pytest==2.3.5
+pytest-xdist==1.8

ar2video.py → libardrone/ar2video.py

@@ -33,22 +33,27 @@
 import paveparser
 import Image
 
-class ARVideo2:
-
-    def __init__(self):
+class ARVideo2(object):
+    def __init__(self, video_pipe = None):
         self.pngsplit = pngsplitter.PNGSplitter(self)
         self.h264 = h264decoder.H264ToPNG(self.pngsplit)
         self.paveparser = paveparser.PaVEParser(self.h264)
-        self.latest_image = Image.new('RGB', (320, 240))
+        self.latest_image = None
+        self.video_pipe = video_pipe
 
     """
     Called by the PNG splitter when there's an image ready
     """
     def image_ready(self, image):
         self.latest_image = image
+        if self.video_pipe:
+            self.video_pipe.send(image)
 
     """
     Guaranteed to return an image as a PIL Image object.
     """
     def get_image(self):
         return self.latest_image
+
+    def write(self, data):
+        self.paveparser.write(data)

arnetwork.py → libardrone/arnetwork.py

@@ -27,6 +27,9 @@
 import socket
 import threading
 import multiprocessing
+import Image
+import numpy as np
+import StringIO
 
 import libardrone
 
@@ -46,15 +49,15 @@ def __init__(self, nav_pipe, video_pipe, com_pipe, is_ar_drone_2):
         self.is_ar_drone_2 = is_ar_drone_2
         if is_ar_drone_2:
             import ar2video
-            self.ar2video = ar2video.ARVideo2()
+            self.ar2video = ar2video.ARVideo2(self.video_pipe)
         else:
             import arvideo
 
     def run(self):
-        if is_ar_drone_2:
+        if self.is_ar_drone_2:
             video_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            video_socket.setblocking(0)
             video_socket.connect(('192.168.1.1', libardrone.ARDRONE_VIDEO_PORT))
+            video_socket.setblocking(0)
         else:
             video_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
             video_socket.setblocking(0)
@@ -80,16 +83,10 @@ def run(self):
                             break
                     if self.is_ar_drone_2:
                         self.ar2video.write(data)
-                        # For now, immediately decode the image for compatibility
-                        # and scale down to 320x240 for the same reason
-                        image = self.ar2video.get_image()
-                        if image == None:
-                            image = Image(320, 240)
-                        image.thumbnail(320, 240)
-                        image = image.tostring()
+                        # Sending is taken care of by the decoder
                     else:
                         w, h, image, t = arvideo.read_picture(data)
-                    self.video_pipe.send(image)
+                        self.video_pipe.send(image)
                 elif i == nav_socket:
                     while 1:
                         try:
@@ -127,7 +124,9 @@ def run(self):
                 if i == self.drone.video_pipe:
                     while self.drone.video_pipe.poll():
                         image = self.drone.video_pipe.recv()
-                    self.drone.image = image
+                    # Convert image to numpy array
+                    self.drone.image = np.array(Image.open(
+                        StringIO.StringIO(image)))
                 elif i == self.drone.nav_pipe:
                     while self.drone.nav_pipe.poll():
                         navdata = self.drone.nav_pipe.recv()

h264decoder.py → libardrone/h264decoder.py

@@ -36,25 +36,24 @@
 ON_POSIX = 'posix' in sys.builtin_module_names
 
 def enqueue_output(out, queue, outfileobject):
-    for d in iter(out.read, b''):
-        outfileobject.write(d)
-    out.close()
+    while True:
+        r = out.read(100)
+        outfileobject.write(r)
 
 """
 Usage: pass a listener, with a method 'data_ready' which will be called whenever there's output
 from ffmpeg. This will be called in an arbitrary thread. You can later call H264ToPng.get_data_if_any to retrieve
 said data.
 You should then call write repeatedly to write some encoded H.264 data.
 """
-class H264ToPNG:
-
+class H264ToPNG(object):
 	def __init__(self, outfileobject):
-		p = Popen(["ffmpeg", "-i", "-", "-f", "image2pipe", "-vcodec", "png", "-r", "5", "-"], stdin=PIPE, stdout=PIPE)
+		p = Popen(["ffmpeg", "-i", "-", "-f", "image2pipe", "-vcodec", "png", "-r", "5", "-"], stdin=PIPE, stdout=PIPE, stderr=open('/dev/null', 'w'))
 		self.writefd = p.stdin
 		self.q = Queue()
-		t = Thread(target=enqueue_output, args=(p.stdout, q, outfileobject))
+		t = Thread(target=enqueue_output, args=(p.stdout, self.q, outfileobject))
 		t.daemon = True # thread dies with the program
 		t.start()
 
 	def write(self, data):
-		self.writefd.write(data)
+            self.writefd.write(data)

paveparser.py → libardrone/paveparser.py

@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
+import struct
 
 """
 The AR Drone 2.0 allows a tcp client to receive H264 (MPEG4.10 AVC) video
@@ -30,40 +31,54 @@
 """
 Usage: Pass in an output file object into the constructor, then call write on this.
 """
-class PaVEParser:
+class PaVEParser(object):
 
     HEADER_SIZE_SHORT = 64; # sometimes header is longer
 
     def __init__(self, outfileobject):
         self.buffer = ""
-        self.state = handle_header
+        self.state = self.handle_header
         self.outfileobject = outfileobject
+        self.misaligned_frames = 0
+        self.payloads = 0
 
     def write(self, data):
         self.buffer += data
 
         while True:
-            made_progress = self.state(self)
+            made_progress = self.state()
             if not made_progress:
                 return
 
     def handle_header(self):
-        if self.fewer_remaining_than(HEADER_SIZE_SHORT):
+        if self.fewer_remaining_than(self.HEADER_SIZE_SHORT):
             return False
-        (signature, version, video_codec, header_size, self.payload_size, encoded_stream_width, encoded_stream_height, display_width, display_height, frame_number, timestamp, total_chunks, chunk_index, frame_type, control, stream_byte_position_lw, stream_byte_position_uw, stream_id, total_slices, slice_index, header1_size, header2_size, reserved2, advertised_size, reserved3) = struct.unpack(self.buffer.slice(0, HEADER_SIZE_SHORT), "<4sBBHIHHHHIIBBBBIIHBBBB2sI12s")
+        (signature, version, video_codec, header_size, self.payload_size, encoded_stream_width, encoded_stream_height, display_width, display_height, frame_number, timestamp, total_chunks, chunk_index, frame_type, control, stream_byte_position_lw, stream_byte_position_uw, stream_id, total_slices, slice_index, header1_size, header2_size, reserved2, advertised_size, reserved3) = struct.unpack("<4sBBHIHHHHIIBBBBIIHBBBB2sI12s", self.buffer[0:self.HEADER_SIZE_SHORT])
         if signature != "PaVE":
-            raise Exception("Invalid signature: "+signature)
-        self.buffer = self.buffer.slice(header_size)
-        self.state = handle_payload
+            self.state = self.handle_misalignment
+            return True
+        self.buffer = self.buffer[header_size:]
+        self.state = self.handle_payload
+        return True
+
+    def handle_misalignment(self):
+        """Sometimes we start of in the middle of frame - look for the PaVE header."""
+        index = self.buffer.find('PaVE')
+        if index == -1:
+            return False
+        self.misaligned_frames += 1
+        self.buffer = self.buffer[index:]
+        self.state = self.handle_header
         return True
 
     def handle_payload(self):
         if self.fewer_remaining_than(self.payload_size):
              return False
-        self.state = handle_header
-        self.outfileobject.write(self.buffer.slice(0, self.payload_size))
-        self.buffer = self.buffer.slice(self.payload_size)
+        self.state = self.handle_header
+        self.outfileobject.write(self.buffer[0:self.payload_size])
+        self.buffer = self.buffer[self.payload_size:]
+        self.payloads += 1
         return True
 
     def fewer_remaining_than(self, desired_size):
-        return self.buffer.length() < desired_size
+        return len(self.buffer) < desired_size

pngsplitter.py → libardrone/pngsplitter.py

@@ -32,7 +32,7 @@
 """
 Usage: Call put_data repeatedly. An array of PNG files will be returned each time you call it.
 """
-class PNGSplitter:
+class PNGSplitter(object):
 
     def __init__(self, listener):
         self.buffer = ""
@@ -49,14 +49,14 @@ def write(self, data):
         self.buffer += data
 
         while True:
-            (found_png, made_progress) = self.state(self)
+            (found_png, made_progress) = self.state()
             if found_png:
-                listener.image_ready(Image.open(StringIO.StringIO(self.buffer.slice(0, self.offset))))
-                self.buffer = self.buffer.slice(self.offset, self.buffer.length() - self.offset) 
+                self.listener.image_ready(self.buffer[0:self.offset])
+                self.buffer = self.buffer[self.offset:] 
                 self.offset = 0
                 self.state = self.handle_header
             if not made_progress:
-                return results
+                return
 
     def handle_header(self):
         self.pngStartOffset = self.offset
@@ -70,7 +70,7 @@ def handle_chunk_header(self):
         if self.fewer_remaining_than(8):
              return (False, False)
         self.state = self.handle_chunk_data
-        self.chunk = struct.unpack(self.buffer.slice(self.offset, 8), ">I4s")
+        self.chunk = struct.unpack( ">I4s", self.buffer[self.offset:(self.offset + 8)])
         self.offset += 8
         return (False, True)
 
@@ -82,8 +82,8 @@ def handle_chunk_data(self):
         if self.chunk[1] == "IEND":
             return (True, True)
         else:
-            self.state = handle_chunk_header
+            self.state = self.handle_chunk_header
             return (False, True)
 
     def fewer_remaining_than(self, desired_size):
-        return self.buffer.length() < self.offset + desired_size
+        return len(self.buffer) < self.offset + desired_size

libardrone/test_h264_decoder.py

@@ -0,0 +1,17 @@
+import paveparser
+import mock
+import h264decoder
+import os
+
+
+def test_h264_decoder():
+    pngstream = mock.Mock()
+    decoder = h264decoder.H264ToPNG(pngstream)
+    example_video_stream = open(os.path.join(os.path.dirname(__file__), 'paveparser.output'))
+    while True:
+        data = example_video_stream.read(1000)
+        if len(data) == 0:
+            break
+        decoder.write(data)
+
+    assert pngstream.write.called

libardrone/test_paveparser.py

@@ -0,0 +1,17 @@
+import paveparser
+import mock
+import os
+
+
+def test_misalignment():
+    outfile = mock.Mock()
+    p = paveparser.PaVEParser(outfile)
+    example_video_stream = open(os.path.join(os.path.dirname(__file__), 'ardrone2_video_example.capture'))
+    while True:
+        data = example_video_stream.read(1000000)
+        if len(data) == 0:
+            break
+        p.write(data)
+
+    assert outfile.write.called
+    assert p.misaligned_frames < 3

libardrone/test_pngsplitter.py

@@ -0,0 +1,17 @@
+import paveparser
+import mock
+import pngsplitter
+import os
+
+
+def test_pngsplitter():
+    listener = mock.Mock()
+    splitter = pngsplitter.PNGSplitter(listener)
+    example_png_stream = open(os.path.join(os.path.dirname(__file__), 'pngstream.example'))
+    while True:
+        data = example_png_stream.read(1000)
+        if len(data) == 0:
+            break
+        splitter.write(data)
+
+    assert listener.image_ready.call_count > 60

requirements.txt

@@ -0,0 +1,3 @@
+PIL==1.1.7
+mock==1.0.1
+numpy==1.7.1