diff --git a/ogr_io.py b/ogr_io.py
index 54a71d38487c22b9be5be14cb7b704e0b729a994..596aa3146bf450a6b2a337d21bdd089478f6a910 100644
--- a/ogr_io.py
+++ b/ogr_io.py
@@ -2,6 +2,7 @@
 Read/Write OGR functions
 """
 from osgeo import ogr
+import sys
 
 def openToRead(shapefile):
     driver = ogr.GetDriverByName("SQLite")
diff --git a/step4_filter_shapefile.ipynb b/step4_filter_shapefile.ipynb
index 7f2feded5b70510f24504936ab20558ec528d576..1f0bd0a04bea88cb027c8e2a79971436502f3470 100644
--- a/step4_filter_shapefile.ipynb
+++ b/step4_filter_shapefile.ipynb
@@ -31,6 +31,7 @@
    "source": [
     "import os\n",
     "import subprocess\n",
+    "from osgeo import ogr, osr\n",
     "\n",
     "import ogr_io\n",
     "\n",
@@ -41,12 +42,23 @@
     "OUTPUT_DIR = \"output\"\n",
     "\n",
     "# Input / Output filenames\n",
+    "envelope_fname = os.path.join(DATA_DIR, \"morbihan.sqlite\")\n",
     "ndwi_thres_fname = os.path.join(OUTPUT_DIR, \"ndwi_threshold30.tif\")\n",
     "watermask_fname = os.path.join(OUTPUT_DIR, \"watermask.sqlite\")\n",
-    "envelope_fname = os.path.join(OUTPUT_DIR, \"morbihan.sqlite\")\n",
     "results_fname = os.path.join(OUTPUT_DIR, \"results.sqlite\")\n",
     "geojson_fname = os.path.join(OUTPUT_DIR, \"geojson.json\")\n",
     "\n",
+    "# Get SQLite driver\n",
+    "outDriver = ogr.GetDriverByName(\"SQLite\")\n",
+    "\n",
+    "# Remove results if it already exists\n",
+    "if os.path.exists(results_fname):\n",
+    "    outDriver.DeleteDataSource(results_fname)\n",
+    "    \n",
+    "# Remove watermask if it already exists\n",
+    "if os.path.exists(watermask_fname):\n",
+    "    outDriver.DeleteDataSource(watermask_fname)\n",
+    "\n",
     "# Convert the TIF file in a shapefile with polygons\n",
     "try:\n",
     "    subprocess.call([\"gdal_polygonize.py\", ndwi_thres_fname,\"-f\", \"SQLite\", watermask_fname])\n",
@@ -82,15 +94,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from osgeo import ogr, osr\n",
-    "\n",
-    "# Save extent to a new Shapefile\n",
-    "outDriver = ogr.GetDriverByName(\"SQLite\")\n",
-    "\n",
-    "# Remove output shapefile if it already exists\n",
-    "if os.path.exists(results_fname):\n",
-    "    outDriver.DeleteDataSource(results_fname)\n",
-    "\n",
     "# create the spatial reference, WGS84\n",
     "srs = osr.SpatialReference()\n",
     "srs.ImportFromEPSG(32630)\n",
@@ -163,6 +166,13 @@
     "m, dc = display_api.rasters_on_map([raster], OUTPUT_DIR, [DATE], geojson_data=fc)\n",
     "m"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {