Refactoring: extract method

pyshared/jbcr2.py

@@ -18,6 +18,10 @@
 # You should have received a copy of the GNU General Public License
 # along with CR2_Scripts. If not, see <http:https://www.gnu.org/licenses/>.
 
+import os
+import tempfile
+import commands
+import numpy as np
 
 ## replace data strips for specified IFD
 
@@ -41,3 +45,48 @@ def replace_ifd(tiff, k, data):
  else:
  raise AssertionError("Reference to data strip not found in IFD#%d" % k)
  return
+
+## convert lossless JPEG encoded input file to raw data
+
+def decode_lossless_jpeg(filename):
+ # create folder for temporary files
+ tmpfolder = tempfile.mkdtemp()
+
+ # decode input file with Stanford PVRG software
+ cmd = 'pvrg-jpeg -d -s "%s" -o "%s"' % (filename, os.path.join(tmpfolder, "parts"))
+ st, out = commands.getstatusoutput(cmd)
+ if st != 0:
+ raise AssertionError('Error decoding JPEG file: %s' % out)
+
+ # interpret output to determine the number of color components and precision
+ components = []
+ for line in out.split('\n'):
+ if line.startswith('>> '):
+ record = line.split()
+ f = record[4]
+ w = int(record[6])
+ h = int(record[8])
+ components.append((f,w,h))
+ elif line.startswith('Caution: precision type:'):
+ record = line.split()
+ print "RAW data precision: %d" % int(record[3])
+ # number of color components
+ n = len(components)
+ # first assemble color components
+ width = sum([w for f,w,h in components])
+ height = components[0][2]
+ assert all([h == height for f,w,h in components])
+ a = np.zeros((height, width), dtype=np.dtype('>H'))
+ for i, (f,w,h) in enumerate(components):
+ # read raw data for this component
+ b = np.fromfile(f, dtype=np.dtype('>H'))
+ b.shape = (h,w)
+ # insert into assembled color image
+ a[:,i::n] = b
+ # remove temporary file
+ os.remove(f)
+
+ # remove (empty) temporary folder
+ os.rmdir(tmpfolder)
+
+ return a, len(components)

raw_decode.py

@@ -21,12 +21,12 @@
 import sys
 import os
 import argparse
-import commands
 import numpy as np
 import matplotlib.pyplot as plt
 
 sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)),'pyshared'))
 import jbtiff
+import jbcr2
 
 ## main program
 
@@ -52,44 +52,16 @@ def main():
  # each slice takes: 432 pixels from each of 4 colors (first two)
  # 472 pixels from each of 4 colors (last one)
 
- # obtain required parameters from RAW file
+ # read input raw file
  tiff = jbtiff.tiff_file(open(args.raw, 'rb'))
- width,height = tiff.get_sensor_size()
- slices = tiff.get_slices()
 
  # convert lossless JPEG encoded input file to raw data
- cmd = 'pvrg-jpeg -d -s "%s" -o parts' % args.input
- st, out = commands.getstatusoutput(cmd)
- if st != 0:
- raise AssertionError('Error decoding JPEG file: %s' % out)
-
- # interpret output to determine color components and precision
- components = []
- for line in out.split('\n'):
- if line.startswith('>> '):
- record = line.split()
- f = record[4]
- w = int(record[6])
- h = int(record[8])
- components.append((f,w,h))
- elif line.startswith('Caution: precision type:'):
- record = line.split()
- print "RAW data precision: %d" % int(record[3])
- # number of color components
- n = len(components)
- # first assemble color components
- assert all([h == height for f,w,h in components])
- assert sum([w for f,w,h in components]) == width
- a = np.zeros((height, width), dtype=np.dtype('>H'))
- for i, (f,w,h) in enumerate(components):
- # read raw data for this color component
- b = np.fromfile(f, dtype=np.dtype('>H'))
- b.shape = (h,w)
- # insert into assembled color image
- a[:,i::n] = b
- # remove temporary file
- os.remove(f)
+ a, components = jbcr2.decode_lossless_jpeg(args.input)
+ height,width = a.shape
 
+ # obtain required parameters from RAW file
+ assert height,width == tiff.get_sensor_size()
+ slices = tiff.get_slices()
  # make a list of the width of each slice
  slice_widths = [slices[1]] * slices[0] + [slices[2]]
  assert sum(slice_widths) == width