home *** CD-ROM | disk | FTP | other *** search
/ MacHack 2000 / MacHack 2000.toast / pc / The Hacks / MacHacksBug / Python 1.5.2c1 / Extensions / img / Lib / imgbmp.py < prev    next >
Encoding:
Python Source  |  2000-06-23  |  5.4 KB  |  211 lines
  1. "Template for image file reader/writers in Python"
  2. #
  3. # Windows BMP image reader/writer module, python version
  4. #
  5. import struct
  6. import imgformat
  7. import imgcolormap
  8. import os
  9.  
  10. # test little/bigendian
  11. _little_endian = (struct.unpack('h', '\001\000')[0] & 0xff)
  12.  
  13. _FILEHEADERSIZE=14
  14. _IMAGEHEADERSIZE=40
  15.  
  16. error = imgformat.error
  17.  
  18. class reader:
  19.     "Object that reads the image."
  20.     
  21.     def __init__(self, file):
  22.         """Initialize. You should read the header and fill in attributes
  23.         such as width, height, format_choices, format and colormap"""
  24.         
  25.         if type(file) == type(''):
  26.             self._filename = filename
  27.             self._fp = open(filename, 'rb')
  28.         else:
  29.             self._filename = "<open file>"
  30.             self._fp = file
  31.         self._readfileheader()
  32.         self._readimageheader()
  33.         if self.bpp == 8:
  34.             self._readcolormap()
  35.             self.format = imgformat.colormap_b2t
  36.         elif self.bpp == 24:
  37.             if _little_endian:
  38.                 self.format = imgformat.bmprgble_b2t
  39.             else:
  40.                 self.format = imgformat.bmprgbbe_b2t
  41.         self.format_choices = (self.format, )
  42.         self._fp.seek(self._dataoffset)
  43.         if self._datasize == 0:
  44.             rowlen = self.width*self.bpp/8
  45.             rowlen = (rowlen+3) & ~3
  46.             self._datasize = rowlen*self.height
  47. ##             print 'DBG: computed datasize', self._datasize
  48.             
  49.     def _readfileheader(self):
  50.         d = self._fp.read(2)
  51.         if d <> 'BM':
  52.             raise error, 'Not a BMP file'
  53.         d = self._fp.read(8)
  54.         self._dataoffset = self._readlong()
  55.         
  56.     def _readimageheader(self):
  57.         wantedhsize = 24
  58.         hsize = self._readlong()
  59.         if hsize < wantedhsize:
  60.             raise error, ('Unexpected header size:', hsize)
  61.         self.width = self._readlong()
  62.         self.height = self._readlong()
  63.         self._fp.read(2) # Skip Numberofimageplanes
  64.         self.bpp = self._readshort()
  65.         if self.bpp not in (8, 24):
  66.             raise error, ('Unsupported number of bits per pixel:', self.bpp)
  67.         cmethod = self._readlong()
  68.         if cmethod <> 0:
  69.             raise error, ('Unsupported compression method:', cmethod)
  70.         self._datasize = self._readlong()
  71.         self._fp.read(hsize-wantedhsize)
  72.         
  73.     def _readcolormap(self):
  74.         ncolors = 1 << self.bpp
  75.         r = g = b = ''
  76.         for i in range(ncolors):
  77.             b = b + self._fp.read(1)
  78.             g = g + self._fp.read(1)
  79.             r = r + self._fp.read(1)
  80.             self._fp.read(1)
  81.         self.colormap = imgcolormap.new3(r, g, b)
  82.         
  83.     def read(self):
  84.         """Read the image data"""
  85.         rv = self._fp.read(self._datasize)
  86.         if len(rv) <> self._datasize:
  87.             raise error, 'Incorrect size read: expected %d got %d'%(self._datasize, len(rv))
  88.         return rv
  89.             
  90.     def _readlong(self):
  91.         """Read a little-endian long, assumed to be positive"""
  92.         d = ord(self._fp.read(1)) & 0xff
  93.         d = d | ((ord(self._fp.read(1)) & 0xff) << 8)
  94.         d = d | ((ord(self._fp.read(1)) & 0xff) << 16)
  95.         d = d | ((ord(self._fp.read(1)) & 0xff) << 24)
  96.         return d
  97.         
  98.     def _readshort(self):
  99.         """Read a little-endian short, assumed to be positive"""
  100.         d = ord(self._fp.read(1)) & 0xff
  101.         d = d | ((ord(self._fp.read(1)) & 0xff) << 8)
  102.         return d
  103.  
  104.     def args(self):
  105.         rv = {}
  106.         for k in self.__dict__.keys():
  107.             if k[0] <> '_':
  108.                 rv[k] = self.__dict__[k]
  109.         return rv
  110.         
  111.  
  112.     def write(self, data):
  113.         raise error, 'Cannot write() to reader'
  114.  
  115. class writer:
  116.     "Object that writes to an image file"
  117.     
  118.     def __init__(self, file):
  119.         if type(file) == type(''):
  120.             self._filename = filename
  121.             self._fp = None
  122.         else:
  123.             self._filename = "<open file>"
  124.             self._fp = file
  125.         if _little_endian:
  126.             rgbformat = imgformat.bmprgble_b2t
  127.         else:
  128.             rgbformat = imgformat.bmprgbbe_b2t
  129.         self.format = imgformat.colormap_b2t
  130.         self.format_choices = (self.format, rgbformat)
  131.  
  132.     def args(self):
  133.         return self.__dict__
  134.         
  135.     def _get(self, attr):
  136.         try:
  137.             return getattr(self, attr)
  138.         except AttributeError:
  139.             raise error, "Required attribute '%s' missing"%attr
  140.  
  141.     def read(self):
  142.         raise error, 'Cannot read() from writer'
  143.  
  144.     def write(self, data):
  145.         """Write the image file, according to attribute format"""
  146.         
  147.         w = self._get('width')
  148.         h = self._get('height')
  149.         if self.format == imgformat.colormap_b2t:
  150.             colormap = self._get('colormap')
  151.             self._bpp = 8
  152.         else:
  153.             colormap = None
  154.             self._bpp = 24
  155.             
  156.         if not self._fp:
  157.             self._fp = open(self._filename, 'wb')
  158.             if os.name == 'mac':
  159.                 import macfs
  160.                 fss = macfs.FSSpec(self._filename)
  161.                 fss.SetCreatorType('????', 'BMP ')
  162.             
  163.         if colormap:
  164.             self._colormapsize = len(colormap)
  165.         else:
  166.             self._colormapsize = 0
  167.         self._imgoffset = _FILEHEADERSIZE + _IMAGEHEADERSIZE + \
  168.                 self._colormapsize*4
  169.         self._datasize = len(data)
  170.         self._writefileheader()
  171.         self._writeimageheader()
  172.         if colormap:
  173.             self._writecolormap(colormap)
  174.         self._fp.write(data)
  175.         self._fp.close()
  176.  
  177.     def _writefileheader(self):            
  178.         self._fp.write('BM')
  179.         self._writelong(self._imgoffset + self._datasize)
  180.         self._writeshort(0)
  181.         self._writeshort(0)
  182.         self._writelong(self._imgoffset)
  183.         
  184.     def _writeimageheader(self):
  185.         self._writelong(40)
  186.         self._writelong(self.width)
  187.         self._writelong(self.height)
  188.         self._writeshort(1)    # Number of image planes
  189.         self._writeshort(self._bpp)
  190.         self._writelong(0)    # compression method
  191.         self._writelong(self._datasize)
  192.         self._writelong(0)    # Horizontal resolution
  193.         self._writelong(0)    # Vertical resolution
  194.         self._writelong(self._colormapsize)
  195.         self._writelong(0)    # Number of "significant" colors
  196.         
  197.     def _writecolormap(self, colormap):
  198.         for r, g, b in colormap:
  199. ##             print 'DBG map', r, g, b
  200.             self._fp.write(chr(b)+chr(g)+chr(r)+chr(0))
  201.         
  202.     def _writeshort(self, d):
  203.         self._fp.write(chr(d & 0xff))
  204.         self._fp.write(chr((d>>8)&0xff))
  205.         
  206.     def _writelong(self, d):
  207.         self._fp.write(chr(d & 0xff))
  208.         self._fp.write(chr((d>>8)&0xff))
  209.         self._fp.write(chr((d>>16)&0xff))
  210.         self._fp.write(chr((d>>24)&0xff))
  211.