Added ability to wrap more than one codestream. #206

glymur/jp2k.py

@@ -18,6 +18,7 @@ else:
 
 from collections import Counter
 import ctypes
+import itertools
 import math
 import os
 import re
@@ -606,7 +607,7 @@ class Jp2k(Jp2kBox):
         self.parse()
 
     def wrap(self, filename, boxes=None):
-        """Write the codestream back out to file, wrapped in new JP2 jacket.
+        """Create a new JP2/JPX file wrapped in a new jacket.
 
         Parameters
         ----------
@@ -640,10 +641,22 @@ class Jp2k(Jp2kBox):
             height = codestream.segment[1].ysiz
             width = codestream.segment[1].xsiz
             num_components = len(codestream.segment[1].xrsiz)
+            if num_components < 3:
+                colorspace = GREYSCALE
+            else:
+                if len(self.box) == 0:
+                    # Best guess is SRGB
+                    colorspace = SRGB
+                else:
+                    # Take whatever the first jp2 header / color specification
+                    # says.
+                    jp2hs = [box for box in self.box if box.box_id == 'jp2h']
+                    colorspace = jp2hs[0].box[1].colorspace
+
             boxes[2].box = [ImageHeaderBox(height=height,
                                            width=width,
                                            num_components=num_components),
-                            ColourSpecificationBox(colorspace=SRGB)]
+                            ColourSpecificationBox(colorspace=colorspace)]
 
         _validate_jp2_box_sequence(boxes)
 
@@ -652,28 +665,30 @@ class Jp2k(Jp2kBox):
                 if box.box_id != 'jp2c':
                     box.write(ofile)
                 else:
-                    # The codestream gets written last.
+                    # Codestreams require a bit more care.
                     if len(self.box) == 0:
                         # Am I a raw codestream?  If so, then it is pretty
                         # easy, just write the codestream box header plus all
                         # of myself out to file.
                         ofile.write(struct.pack('>I', self.length + 8))
-                        ofile.write('jp2c'.encode())
+                        ofile.write(b'jp2c')
                         with open(self.filename, 'rb') as ifile:
                             ofile.write(ifile.read())
                     else:
                         # OK, I'm a jp2 file.  Need to find out where the
                         # raw codestream actually starts.
-                        jp2c = [box for box in self.box
-                                if box.box_id == 'jp2c']
-                        jp2c = jp2c[0]
-                        ofile.write(struct.pack('>I', jp2c.length))
-                        ofile.write('jp2c'.encode())
+                        offset = box.offset
+                        length = box.length
+                        if offset == -1:
+                            # Find the first codestream in the file.
+                            jp2c = [box for box in self.box
+                                    if box.box_id == 'jp2c']
+                            offset = jp2c[0].offset
+                            length = jp2c[0].length
+
                         with open(self.filename, 'rb') as ifile:
-                            # Seek 8 bytes past the L, T fields to get to the
-                            # raw codestream.
-                            ifile.seek(jp2c.offset + 8)
-                            ofile.write(ifile.read(jp2c.length - 8))
+                            ifile.seek(offset)
+                            ofile.write(ifile.read(length))
 
             ofile.flush()
 
@@ -1301,7 +1316,7 @@ def _check_jp2h_child_boxes(boxes, parent_box_name):
     """Certain boxes can only reside in the JP2 header."""
     box_ids = set([box.box_id for box in boxes])
     intersection = box_ids.intersection(JP2H_CHILDREN)
-    if len(intersection) > 0 and parent_box_name != 'jp2h':
+    if len(intersection) > 0 and parent_box_name not in ['jp2h', 'jpch']:
         msg = "A '{0}' box can only be nested in a JP2 header box."
         raise IOError(msg.format(list(intersection)[0]))
 

glymur/test/test_jp2box.py

@@ -58,6 +58,23 @@ class TestDataEntryURL(unittest.TestCase):
     def setUp(self):
         self.jp2file = glymur.data.nemo()
 
+    def test_wrap_greyscale(self):
+        """A single component should be wrapped as GREYSCALE."""
+        j = Jp2k(self.jp2file)
+        data = j.read()
+        red = data[:, :, 0]
+
+        # Write it back out as a raw codestream.
+        with tempfile.NamedTemporaryFile(suffix=".j2k") as tfile1:
+            j2k = glymur.Jp2k(tfile1.name, 'wb')
+            j2k.write(data[:, :, 0])
+
+            # Ok, now rewrap it as JP2.  The colorspace should be GREYSCALE.
+            with tempfile.NamedTemporaryFile(suffix=".jp2") as tfile2:
+                jp2 = j2k.wrap(tfile2.name)
+                self.assertEqual(jp2.box[2].box[1].colorspace,
+                        glymur.core.GREYSCALE)
+
     def test_basic_url(self):
         """Just your most basic URL box."""
         # Wrap our j2k file in a JP2 box along with an interior url box.

glymur/test/test_jp2box_jpx.py

@@ -23,6 +23,7 @@ class TestJPXWrap(unittest.TestCase):
     """Test suite for wrapping JPX files."""
 
     def setUp(self):
+        self.jpxfile = glymur.data.jpxfile()
         self.jp2file = glymur.data.nemo()
         self.j2kfile = glymur.data.goodstuff()
 
@@ -44,6 +45,17 @@ class TestJPXWrap(unittest.TestCase):
     def tearDown(self):
         os.unlink(self.xmlfile)
 
+    def test_full_blown_jpx(self):
+        """Rewrap a jpx file."""
+        with tempfile.NamedTemporaryFile(suffix='.jp2') as tfile1:
+            jpx = Jp2k(self.jpxfile)
+            idx = list(range(5)) + list(range(6, 9)) + list(range(9, 12)) + [12]
+            boxes = [jpx.box[j] for j in idx]
+            jpx2 = jpx.wrap(tfile1.name, boxes=boxes)
+            exp_ids = [box.box_id for box in boxes]
+        act_ids = [box.box_id for box in jpx2.box]
+        self.assertEqual(exp_ids, act_ids)
+
     def test_jpx_ftbl_no_codestream(self):
         """Can have a jpx with no codestream."""
         with tempfile.NamedTemporaryFile(suffix='.jp2') as tfile1: