Skip to content

Commit

Permalink
Add image context in Jupyter notebook to be able to restart a step
Browse files Browse the repository at this point in the history
  • Loading branch information
sbrunner committed Aug 14, 2023
1 parent 7ad3470 commit 40d3ac4
Showing 1 changed file with 54 additions and 20 deletions.
74 changes: 54 additions & 20 deletions scan_to_paperless/process.py
Original file line number Diff line number Diff line change
Expand Up @@ -1521,7 +1521,10 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
# Open one of the images
context.image = cv2.imread(os.path.join(base_folder, "{step["sources"][0]}"))
{other_images_open}"""
{other_images_open}
images_context = {{"original": context.image.clone()}}
"""
)
)

Expand Down Expand Up @@ -1554,7 +1557,7 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
np.arange(0, context.image.shape[2]),
)
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))"""
display(Image.fromarray(cv2.cvtColor(images_context["original"][context.index], cv2.COLOR_BGR2RGB)))"""
)
)

Expand Down Expand Up @@ -1582,8 +1585,9 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""context.config["args"]["auto_mask"] = {_pretty_repr(context.config["args"].get("auto_mask", {}), " ")}
f"""context.image = images_context["original"]
context.config["args"]["auto_mask"] = {_pretty_repr(context.config["args"].get("auto_mask", {}), " ")}
hsv = cv2.cvtColor(context.image, cv2.COLOR_BGR2HSV)
print("Hue (h)")
Expand All @@ -1607,7 +1611,10 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
context.init_mask()
if context.mask is not None:
display(Image.fromarray(cv2.cvtColor(context.mask, cv2.COLOR_GRAY2RGB)[context.index]))
display(Image.fromarray(cv2.cvtColor(context.get_masked()[context.index], cv2.COLOR_BGR2RGB)))"""
display(Image.fromarray(cv2.cvtColor(context.get_masked()[context.index], cv2.COLOR_BGR2RGB)))
images_context["auto_mask"] = context.image.clone()
"""
)
)

Expand All @@ -1616,7 +1623,9 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""context.config["args"]["level"] = {context.config["args"].get("level", schema.LEVEL_DEFAULT)},
f"""context.image = images_context["auto_mask"]
context.config["args"]["level"] = {context.config["args"].get("level", schema.LEVEL_DEFAULT)},
context.config["args"]["cut_white"] = {context.config["args"].get("cut_white", schema.CUT_WHITE_DEFAULT)},
context.config["args"]["cut_black"] = {context.config["args"].get("cut_black", schema.CUT_BLACK_DEFAULT)},
Expand All @@ -1633,11 +1642,15 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""context.config["args"]["auto_level"] = {context.config["args"].get("auto_level", schema.AUTO_LEVEL_DEFAULT)},
f"""context.image = images_context["auto_mask"]
context.config["args"]["auto_level"] = {context.config["args"].get("auto_level", schema.AUTO_LEVEL_DEFAULT)},
context.config["args"]["level"] = {context.config["args"].get("level", schema.LEVEL_DEFAULT)},
process.level(context)
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))"""
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))
images_context["level"] = context.image.clone()"""
)
)

Expand All @@ -1650,12 +1663,15 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""
f"""context.image = images_context["level"]
print(f"Use cut_white: {context.config["args"]["cut_white"]}")
print(f"Use cut_black: {context.config["args"]["cut_black"]}")
process.color_cut(context)
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))"""
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))
images_context["color_cut"] = context.image.clone()"""
)
)

Expand All @@ -1673,15 +1689,19 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""context.config["args"]["auto_cut"] = {_pretty_repr(context.config["args"].get("auto_cut", {}), " ")}"
f"""context.image = images_context["color_cut"]
context.config["args"]["auto_cut"] = {_pretty_repr(context.config["args"].get("auto_cut", {}), " ")}"
# Print in HSV some point of the image
hsv = cv2.cvtColor(context.image, cv2.COLOR_BGR2HSV)
print("Pixel 10:10: ", hsv[10, 10])
print("Pixel 100:100: ", hsv[100, 100])
process.cut(context)
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))"""
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))
images_context["cut"] = context.image.clone()"""
)
)

Expand All @@ -1690,8 +1710,10 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""context.config["args"]["deskew"] = {_pretty_repr(context.config["args"].get("deskew", {}), " ")},
}}
f"""context.image = images_context["cut"]
context.config["args"]["deskew"] = {_pretty_repr(context.config["args"].get("deskew", {}), " ")}
# The angle can be forced in config.images_config.<image_name>.angle.
process.deskew(context)
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))"""
Expand All @@ -1705,10 +1727,14 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""context.config["args"]["crop"] = {_pretty_repr(context.config["args"].get("crop", {}), " ")}
f"""context.image = images_context["deskew"]
context.config["args"]["crop"] = {_pretty_repr(context.config["args"].get("crop", {}), " ")}
process.docrop(context)
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))"""
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))
images_context["crop"] = context.image.clone()"""
)
)

Expand All @@ -1717,10 +1743,14 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""context.config["args"]["sharpen"] = {context.config["args"].get("sharpen", schema.SHARPEN_DEFAULT)}
f"""context.image = images_context["crop"]
context.config["args"]["sharpen"] = {context.config["args"].get("sharpen", schema.SHARPEN_DEFAULT)}
process.sharpen(context)
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))"""
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))
images_context["sharpen"] = context.image.clone()"""
)
)

Expand All @@ -1729,10 +1759,14 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
f"""context.config["args"]["dither"] = {context.config["args"].get("dither", schema.DITHER_DEFAULT)}
f"""context.image = images_context["sharpen"]
context.config["args"]["dither"] = {context.config["args"].get("dither", schema.DITHER_DEFAULT)}
process.dither(context)
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))"""
display(Image.fromarray(cv2.cvtColor(context.image[context.index], cv2.COLOR_BGR2RGB)))
images_context["dither"] = context.image.clone()"""
)
)

Expand All @@ -1745,7 +1779,7 @@ def _create_jupyter_notebook(root_folder: str, context: Context, step: schema.St
)
notebook["cells"].append(
nbformat.v4.new_code_cell( # type: ignore[no-untyped-call]
"""context.config["args"] = {}
"""context.image = images_context["dither"]
try:
process.autorotate(context)
Expand Down

0 comments on commit 40d3ac4

Please sign in to comment.