Merge pull request #8 from MITLibraries/raster

Add support for GeoTIFFs

magma/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+
+
+class UnsupportedFormat(Exception):
+    pass

magma/home/__init__.py

@@ -3,7 +3,8 @@
 import tempfile
 import shutil
 import os
-from zipfile import ZipFile
+from zipfile import BadZipfile
+from contextlib import closing
 try:
     from StringIO import StringIO
 except ImportError:
@@ -13,7 +14,8 @@
                    session,)
 from werkzeug import secure_filename
 
-from magma.upload import process
+from magma.upload import process, datasource
+from magma import UnsupportedFormat
 
 
 home_page = Blueprint('home_page', __name__)
@@ -23,37 +25,30 @@
 def index():
     if request.method == 'POST':
         if 'data' in request.files and request.files['data']:
-            shp = request.files['data']
+            datafile = request.files['data']
         else:
-            flash('The shapefile is required to proceed.', 'danger')
+            flash('A datafile is required to proceed.', 'danger')
             return render_template('index.html')
-
-        if 'metadata' in request.files and request.files['metadata']:
-            fgdc = request.files['metadata']
-        else:
-            fgdc = StringIO('<metadata/>')
+        fgdc = request.files.get('metadata') or StringIO('<metadata/>')
 
         tmp = tempfile.mkdtemp()
-        filename = os.path.join(tmp, secure_filename(shp.filename))
-        shp.save(filename)
-        try:
-            archive = ZipFile(filename, 'r')
-        except:
-            flash('File is not a zip file.', 'danger')
-            return render_template('index.html')
+        filename = os.path.join(tmp, secure_filename(datafile.filename))
+        datafile.save(filename)
 
         try:
-            archive.extractall(tmp)
-        finally:
-            archive.close()
-        try:
-            metadata = process(tmp, fgdc)
+            with closing(datasource(filename)) as ds:
+                metadata = process(ds, fgdc)
+            if request.form['access'] == 'restricted':
+                metadata.set_restricted_access()
+            response = make_response(metadata.write())
+            response.headers['Content-type'] = 'text/xml'
+            response.headers['Content-Disposition'] = \
+                'attachment; filename=fgdc.xml'
+            return response
+        except BadZipfile:
+            flash('Could not open zipfile.', 'danger')
+        except UnsupportedFormat:
+            flash('Uploaded datafile should be either a zipped shapefile or a GeoTIFF.', 'danger')
         finally:
             shutil.rmtree(tmp)
-        if request.form['access'] == 'restricted':
-            metadata.set_restricted_access()
-        response = make_response(metadata.write())
-        response.headers['Content-type'] = 'text/xml'
-        response.headers['Content-Disposition'] = 'attachment; filename=fgdc.xml'
-        return response
     return render_template('index.html')

magma/upload.py

@@ -1,12 +1,15 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import
-from contextlib import closing
 from functools import reduce
 import glob
+import os
+from zipfile import ZipFile
 
-from osgeo import ogr
+from osgeo import ogr, gdal
 from lxml import etree
 
+from magma import UnsupportedFormat
+
 
 class Shapefile(object):
     def __init__(self, file):
@@ -35,6 +38,32 @@ def data_type(self):
         return ("Vector", ogr.GeometryTypeToName(self.layer.GetGeomType()))
 
 
+class GeoTIFF(object):
+    attributes = []
+
+    def __init__(self, file):
+        self.ds = gdal.Open(file)
+
+    @property
+    def name(self):
+        return os.path.basename(self.ds.GetDescription())
+
+    @property
+    def extent(self):
+        xform = self.ds.GetGeoTransform()
+        upper_left = xform[0], xform[3]
+        lower_right = transform_point(xform, self.ds.RasterXSize,
+                                      self.ds.RasterYSize)
+        return upper_left[0], lower_right[0], lower_right[1], upper_left[1]
+
+    @property
+    def data_type(self):
+        return 'Raster', 'Pixel'
+
+    def close(self):
+        self.ds = None
+
+
 parser = etree.XMLParser(remove_blank_text=True)
 
 
@@ -125,19 +154,33 @@ def _get_path(self, path):
         return get_path(path, self.root)
 
 
-def process(data, fgdc):
-    files = glob.glob("%s/*.shp" % data)
+def process(ds, fgdc):
     metadata = FGDC(fgdc)
-    with closing(Shapefile(files[0])) as ds:
-        metadata.set_extent(ds.extent).\
-            set_data_type(ds.data_type).\
-            set_attributes(ds.attributes).\
-            set_name(ds.name).\
-            set_distribution().\
-            set_metadata_contact()
+    metadata.set_extent(ds.extent).\
+        set_data_type(ds.data_type).\
+        set_attributes(ds.attributes).\
+        set_name(ds.name).\
+        set_distribution().\
+        set_metadata_contact()
     return metadata
 
 
+def datasource(datafile):
+    ext = os.path.splitext(datafile)[1]
+    if ext == '.zip':
+        extract_dir = os.path.abspath(os.path.dirname(datafile))
+        archive = ZipFile(datafile, 'r')
+        try:
+            archive.extractall(extract_dir)
+        finally:
+            archive.close()
+        shpfile = glob.glob("%s/*.shp" % extract_dir)[0]
+        return Shapefile(shpfile)
+    elif ext in ('.tif', '.tiff'):
+        return GeoTIFF(datafile)
+    raise UnsupportedFormat
+
+
 def get_path(path, root):
     return reduce(get_or_set, path.split('/'), root)
 
@@ -171,3 +214,9 @@ def get_or_set(root, child):
     <cntvoice>617-258-5598</cntvoice>
     <cntemail>gishelp@mit.edu</cntemail>
   </cntinfo>"""
+
+
+def transform_point(xform, x, y):
+    geo_x = xform[0] + x * xform[1] + y * xform[2]
+    geo_y = xform[3] + x * xform[4] + y * xform[5]
+    return geo_x, geo_y

tests/conftest.py

@@ -27,6 +27,16 @@ def shapefile():
     return _fixtures('SDE_DATA_BD_A8GNS_2003/SDE_DATA_BD_A8GNS_2003.shp')
 
 
+@pytest.fixture
+def geotiff():
+    return _fixtures('grayscale.tif')
+
+
+@pytest.fixture
+def fgdc():
+    return _fixtures('fgdc.xml')
+
+
 def _fixtures(path):
     current_dir = os.path.dirname(os.path.realpath(__file__))
     return os.path.join(current_dir, 'fixtures', path)

tests/test_upload.py

@@ -8,7 +8,7 @@
 
 import pytest
 
-from magma.upload import Shapefile, FGDC
+from magma.upload import Shapefile, FGDC, GeoTIFF
 
 
 @pytest.yield_fixture
@@ -17,6 +17,12 @@ def shp(shapefile):
         yield datasource
 
 
+@pytest.yield_fixture
+def tif(geotiff):
+    with closing(GeoTIFF(geotiff)) as datasource:
+        yield datasource
+
+
 def test_shapefile_has_name(shp):
     assert shp.name == 'SDE_DATA_BD_A8GNS_2003'
 
@@ -34,6 +40,22 @@ def test_shapefile_has_data_type(shp):
     assert shp.data_type == ('Vector', 'Point')
 
 
+def test_raster_has_name(tif):
+    assert tif.name == 'grayscale.tif'
+
+
+def test_raster_has_empty_attributes(tif):
+    assert tif.attributes == []
+
+
+def test_raster_has_extent(tif):
+    assert tif.extent == (-80.0, -60.0, 35.0, 45.0)
+
+
+def test_raster_has_data_type(tif):
+    assert tif.data_type == ('Raster', 'Pixel')
+
+
 def test_sets_extent():
     doc = StringIO('<metadata></metadata>')
     r = FGDC(doc)

tests/test_views.py

@@ -27,10 +27,25 @@ def test_submit_with_shapefile_without_metadata(testapp):
 def test_submit_without_shapefile(testapp):
     form = testapp.get('/').form
     res = form.submit()
-    assert 'The shapefile is required to proceed' in res
+    assert 'A datafile is required to proceed' in res
 
 def test_submit_with_invalid_shapefile(testapp):
     form = testapp.get('/').form
     form['data'] = Upload('tests/fixtures/fgdc.xml')
     res = form.submit()
-    assert 'File is not a zip file' in res
+    assert 'Uploaded datafile should be either a zipped shapefile or a GeoTIFF.' in res
+
+
+def test_submits_geotiff_with_metadata(testapp, geotiff, fgdc):
+    form = testapp.get('/').form
+    form['data'] = Upload(geotiff)
+    form['metadata'] = Upload(fgdc)
+    res = form.submit()
+    assert res.content_type == 'text/xml'
+
+
+def test_submits_geotiff_without_metadata(testapp, geotiff):
+    form = testapp.get('/').form
+    form['data'] = Upload(geotiff)
+    res = form.submit()
+    assert res.content_type == 'text/xml'