Skip to content
Snippets Groups Projects
Commit da4b97c0 authored by onb1259's avatar onb1259
Browse files

fixed typo, added print statement

parent e2317bf0
Branches
No related tags found
No related merge requests found
Show whitespace changes
Inline Side-by-side
Showing
with 4 additions and 3 deletions
extract_figures.ipynb
+ 4
− 3
View file @ da4b97c0
%% Cell type:markdown id:82ab6471 tags: %% Cell type:markdown id:82ab6471 tags:
# Extract figures # Extract figures
This notebook uses a YOLOv8 model trained on five annotated ABO books (1700 pages) to extract figures from them. 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. 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: %% Cell type:code id:e1bcdac6 tags:
``` python ``` python
!pip install pillow !pip install pillow
!pip install zipp !pip install zipp
!pip install ultralytics !pip install ultralytics
import os, shutil import os, shutil
from zipapp import zipfile from zipapp import zipfile
from PIL import Image from PIL import Image
from ultralytics import YOLO from ultralytics import YOLO
from IPython.display import display, FileLink, HTML from IPython.display import display, FileLink, HTML
from iiif_utils import get_imgurls_from_manifesturl, create_paths_from_iiifurls, download_images_multithreded from iiif_utils import get_imgurls_from_manifesturl, create_paths_from_iiifurls, download_images_multithreded
# Url of iiif manifest # Url of iiif manifest
input_manifest = 'https://iiif.onb.ac.at/presentation/ABO/+Z105572305/manifest/' input_manifest = 'https://iiif.onb.ac.at/presentation/ABO/+Z105572305/manifest/'
``` ```
%% Cell type:code id:65afd0ba tags: %% Cell type:code id:65afd0ba tags:
``` python ``` python
# Donwload the book pages # Donwload the book pages
IMAGES_DIR = 'images' IMAGES_DIR = 'images'
imgurls = get_imgurls_from_manifesturl(input_manifest) imgurls = get_imgurls_from_manifesturl(input_manifest)
imgpaths = create_paths_from_iiifurls(imgurls) imgpaths = create_paths_from_iiifurls(imgurls)
if os.path.isdir(IMAGES_DIR): if os.path.isdir(IMAGES_DIR):
shutil.rmtree(IMAGES_DIR) shutil.rmtree(IMAGES_DIR)
os.mkdir(IMAGES_DIR) os.mkdir(IMAGES_DIR)
imgpaths = [os.path.join(IMAGES_DIR, e) for e in imgpaths] imgpaths = [os.path.join(IMAGES_DIR, e) for e in imgpaths]
download_images_multithreded(imgurls, imgpaths, 10) download_images_multithreded(imgurls, imgpaths, 10)
print(f'Total number of downloaded pages: {len(os.listdir(IMAGES_DIR))}')
``` ```
%% Cell type:code id:c6b37e20 tags: %% Cell type:code id:c6b37e20 tags:
``` python ``` python
# Apply figure extraction model # Apply figure extraction model
RESULTS_DIR = 'extracted_figures' RESULTS_DIR = 'extracted_figures'
if os.path.isdir(RESULTS_DIR): if os.path.isdir(RESULTS_DIR):
shutil.rmtree(RESULTS_DIR) shutil.rmtree(RESULTS_DIR)
os.mkdir(RESULTS_DIR) os.mkdir(RESULTS_DIR)
model = YOLO('model_extract_figures.pt') model = YOLO('model_extract_figures.pt')
counter = 0 counter = 0
for path in imgpaths: for path in imgpaths:
result = model.predict(path, verbose=False) result = model.predict(path, verbose=False)
img = Image.open(path) img = Image.open(path)
for i, bb in enumerate(result[0].boxes.xyxy.tolist()): for i, bb in enumerate(result[0].boxes.xyxy.tolist()):
#print(f'Figure detected in: {path}') #print(f'Figure detected in: {path}')
counter += 1 counter += 1
bb = [int(e) for e in bb] bb = [int(e) for e in bb]
img_cropped = img.crop(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])) 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) img_cropped.save(path_cropped)
print(f'Total number of extracted figures: {counter}') print(f'Total number of extracted figures: {counter}')
``` ```
%% Cell type:code id:a997833e tags: %% Cell type:code id:a997833e tags:
``` python ``` python
# Create a zip file with extracted figures and show them # Create a zip file with extracted figures and show them
RESULT_ZIP = 'exatrcted_figures.zip' RESULT_ZIP = 'extracted_figures.zip'
with zipfile.ZipFile(RESULT_ZIP, 'w') as myzip: with zipfile.ZipFile(RESULT_ZIP, 'w') as myzip:
for f in os.listdir(RESULTS_DIR): for f in os.listdir(RESULTS_DIR):
myzip.write(os.path.join(RESULTS_DIR, f)) myzip.write(os.path.join(RESULTS_DIR, f))
zipped_extracted_figures = FileLink(RESULT_ZIP, result_html_prefix="Click here to download zipped extracted figures: ") zipped_extracted_figures = FileLink(RESULT_ZIP, result_html_prefix="Click here to download zipped extracted figures: ")
display(zipped_extracted_figures) display(zipped_extracted_figures)
def show_images(filenames, width, margin): def show_images(filenames, width, margin):
lis = [f'<img src="{name}" style="display:inline;margin:{margin}px" width="{width}"/>' for name in filenames] lis = [f'<img src="{name}" style="display:inline;margin:{margin}px" width="{width}"/>' for name in filenames]
html = ''.join(lis) html = ''.join(lis)
display(HTML(html)) display(HTML(html))
show_images([os.path.join(RESULTS_DIR, e) for e in sorted(os.listdir(RESULTS_DIR))], 150, 1) show_images([os.path.join(RESULTS_DIR, e) for e in sorted(os.listdir(RESULTS_DIR))], 150, 1)
``` ```
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment