fix not existing directories

e2317bf0 · onb1259 · 7e86a177 · e2317bf0
Commit e2317bf0 authored 1 year ago by onb1259
--- a/extract_figures.ipynb
+++ b/extract_figures.ipynb
@@ -48,7 +48,8 @@
    "IMAGES_DIR = 'images'\n",
    "imgurls = get_imgurls_from_manifesturl(input_manifest)\n",
    "imgpaths = create_paths_from_iiifurls(imgurls)\n",
-    "shutil.rmtree(IMAGES_DIR)\n",
+    "if os.path.isdir(IMAGES_DIR):\n",
+    "    shutil.rmtree(IMAGES_DIR)\n",
    "os.mkdir(IMAGES_DIR)\n",
    "imgpaths = [os.path.join(IMAGES_DIR, e) for e in imgpaths]\n",
    "download_images_multithreded(imgurls, imgpaths, 10)"
@@ -66,7 +67,8 @@
    "# Apply figure extraction model\n",
    "\n",
    "RESULTS_DIR = 'extracted_figures'\n",
-    "shutil.rmtree(RESULTS_DIR)\n",
+    "if os.path.isdir(RESULTS_DIR):\n",
+    "    shutil.rmtree(RESULTS_DIR)\n",
    "os.mkdir(RESULTS_DIR)\n",
    "model = YOLO('model_extract_figures.pt')\n",
    "counter = 0\n",

 %% Cell type:markdown id:82ab6471 tags:

 # Extract figures

 This notebook uses a YOLOv8 model trained on five annotated ABO books (1700 pages) to extract figures from them.
 After giving the iiif manifest url of the book from which you want to extract images, the book pages get dwonloaded, the figure recognition model gets applied and the resulting bounding boxes get extracted and can be donwloaded as a zip file.

 %% Cell type:code id:e1bcdac6 tags:

 ``` python
 !pip install pillow
 !pip install zipp
 !pip install ultralytics

 import os, shutil
 from zipapp import zipfile
 from PIL import Image
 from ultralytics import YOLO
 from IPython.display import display, FileLink, HTML

 from iiif_utils import get_imgurls_from_manifesturl, create_paths_from_iiifurls, download_images_multithreded

 # Url of iiif manifest
 input_manifest = 'https://iiif.onb.ac.at/presentation/ABO/+Z105572305/manifest/'
 ```

 %% Cell type:code id:65afd0ba tags:

 ``` python
 # Donwload the book pages

 IMAGES_DIR = 'images'
 imgurls = get_imgurls_from_manifesturl(input_manifest)
 imgpaths = create_paths_from_iiifurls(imgurls)
-shutil.rmtree(IMAGES_DIR)
+if os.path.isdir(IMAGES_DIR):
+    shutil.rmtree(IMAGES_DIR)
 os.mkdir(IMAGES_DIR)
 imgpaths = [os.path.join(IMAGES_DIR, e) for e in imgpaths]
 download_images_multithreded(imgurls, imgpaths, 10)
 ```

 %% Cell type:code id:c6b37e20 tags:

 ``` python
 # Apply figure extraction model

 RESULTS_DIR = 'extracted_figures'
-shutil.rmtree(RESULTS_DIR)
+if os.path.isdir(RESULTS_DIR):
+    shutil.rmtree(RESULTS_DIR)
 os.mkdir(RESULTS_DIR)
 model = YOLO('model_extract_figures.pt')
 counter = 0
 for path in imgpaths:
    result = model.predict(path, verbose=False)
    img = Image.open(path)
    for i, bb in enumerate(result[0].boxes.xyxy.tolist()):
        #print(f'Figure detected in: {path}')
        counter += 1
        bb = [int(e) for e in bb]
        img_cropped = img.crop(bb)
        path_cropped = os.path.join(RESULTS_DIR, os.path.basename(path.rsplit('.', 1)[0] + '_crop' + str(i) + '.' + path.rsplit('.', 1)[1]))
        img_cropped.save(path_cropped)
 print(f'Total number of extracted figures: {counter}')
 ```

 %% Cell type:code id:a997833e tags:

 ``` python
 # Create a zip file with extracted figures and show them

 RESULT_ZIP = 'exatrcted_figures.zip'
 with zipfile.ZipFile(RESULT_ZIP, 'w') as myzip:
    for f in os.listdir(RESULTS_DIR):
        myzip.write(os.path.join(RESULTS_DIR, f))
 zipped_extracted_figures = FileLink(RESULT_ZIP, result_html_prefix="Click here to download zipped extracted figures: ")
 display(zipped_extracted_figures)

 def show_images(filenames, width, margin):
    lis = [f'<img src="{name}" style="display:inline;margin:{margin}px" width="{width}"/>' for name in filenames]
    html = ''.join(lis)
    display(HTML(html))
 show_images([os.path.join(RESULTS_DIR, e) for e in sorted(os.listdir(RESULTS_DIR))], 150, 1)
 ```