Add streamlit app code v1

.gitignore

@@ -0,0 +1,147 @@
+# Cache and other mess
+/temp/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# IDEs
+/.idea/
+/.vscode/
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+/wandb/
+wandb/

.streamlit/config.toml

@@ -0,0 +1,3 @@
+[server]
+
+maxUploadSize = 64

README.md

@@ -1,12 +1,15 @@
 ---
-title: Pxdn Line Extractor
-emoji: 🌖
-colorFrom: blue
-colorTo: red
+title: PXDN Line Extractor
+emoji: ✏️
+colorFrom: pink
+colorTo: indigo
 sdk: streamlit
 sdk_version: 1.41.1
 app_file: app.py
 pinned: false
+license: none
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Line Extractor v1
+
+An attention based convolutional neural network that extracts lineart from a given rgb image.

app.py

@@ -0,0 +1,181 @@
+import io
+import os
+import sys
+
+sys.path.append(os.path.join(os.path.dirname(__file__)))
+
+import blended_tiling
+import numpy
+import onnxruntime
+import streamlit.file_util
+import torch
+import torch.cuda
+from PIL import Image
+from streamlit.runtime.uploaded_file_manager import UploadedFile
+from streamlit_image_comparison import image_comparison
+from torchvision.transforms import functional as TVTF
+
+from tools import image_tools
+
+# * Cached/loaded model
+onnx_session = None  # type: onnxruntime.InferenceSession
+
+# * Streamlit UI / Config
+streamlit.set_page_config(page_title="🐲 PXDN Line Extractor v1", layout="wide")
+streamlit.title("🐲 PXDN Line Extractor v1")
+
+# * Streamlit Containers / Base Layout
+# Row 1
+ui_section_status = streamlit.container()
+
+# Row 2
+ui_col1, ui_col2 = streamlit.columns(2, gap="medium")
+streamlit.html("<hr>")
+
+# Row 3
+ui_section_compare = streamlit.container()
+
+# * Streamlit Session
+# Nothing yet
+
+with ui_section_status:
+	# Forward declared UI elements
+	ui_status_text = streamlit.empty()
+	ui_progress_bar = streamlit.empty()
+
+with ui_col1:
+	# Input Area
+	streamlit.markdown("### Input Image")
+	ui_image_input = streamlit.file_uploader("Upload an image", key="fileupload_image", type=[".png", ".jpg", ".jpeg", ".webp"])  # type: UploadedFile
+
+with ui_col2:
+	# Output Area
+	streamlit.markdown("### Output Image")
+	# Preallocate image spot and download button
+	ui_image_output = streamlit.empty()
+	ui_image_download = streamlit.empty()
+
+def fetch_model_to_cache(huggingface_repo: str, file_path: str, access_token: str) -> str:
+	import huggingface_hub
+	return huggingface_hub.hf_hub_download(huggingface_repo, file_path, token=access_token)
+
+def bootstrap_model():
+	global onnx_session
+	if onnx_session is None:
+		
+		# Environment-level configuration
+		huggingface_repo = os.getenv("HF_REPO_NAME", "")
+		file_path = os.getenv("HF_FILE_PATH", "")
+		access_token = os.getenv("HF_TOKEN", "")
+		allow_cuda = bool(os.getenv("ALLOW_CUDA", "false").lower())
+		
+		model_file_path = fetch_model_to_cache(huggingface_repo, file_path, access_token)
+		
+		# * Enable CUDA if available and allowed
+		model_providers = ['CPUExecutionProvider']
+		if torch.cuda.is_available() and allow_cuda:
+			model_providers.insert(0, 'CUDAExecutionProvider')
+		
+		onnx_session = onnxruntime.InferenceSession(model_file_path, sess_options=None, providers=model_providers)
+
+def evaluate_tiled(image_pt: torch.Tensor, tile_size: int = 128, batch_size: int = 1) -> Image.Image:
+	image_pt_orig = image_pt
+	orig_h, orig_w = image_pt_orig.shape[1], image_pt_orig.shape[2]
+	
+	# ? Padding
+	image_pt_padded, place_x, place_y = image_tools.pad_to_divisible(image_pt_orig, tile_size)
+	
+	_, im_h_padded, im_w_padded = image_pt_padded.shape
+	
+	# ? Tiling
+	image_tiler = blended_tiling.TilingModule(tile_size=tile_size, tile_overlap=[0.18, 0.18], base_size=(im_w_padded, im_h_padded)).eval()
+	# * Add batch dim for the tiler which expects (1, C, H, W)
+	image_tiles = image_tiler.split_into_tiles(image_pt_padded.unsqueeze(0))
+	
+	# ? Pull the input and output names from the model so we're not hardcoding them.
+	onnx_session.get_modelmeta()
+	input_name = onnx_session.get_inputs()[0].name
+	output_name = onnx_session.get_outputs()[0].name
+	
+	# ? Inference ==================================================================================================
+	complete_tiles = []
+	
+	max_evals = image_tiles.size(0) // batch_size
+	image_tiles = image_tiles.numpy()
+	
+	ui_status_text.markdown("### Processing...")
+	active_progress = ui_progress_bar.progress(0, "Progress")
+	
+	for i in range(max_evals):
+		tile_batch = image_tiles[i * batch_size:(i + 1) * batch_size]
+		if len(tile_batch) == 0:
+			break
+		
+		pct_complete = round((i + 1) / max_evals, 2)
+		active_progress.progress(pct_complete)
+		
+		eval_output = onnx_session.run([output_name], {input_name: tile_batch})
+		output_batch = eval_output[0]
+		
+		complete_tiles.extend(output_batch)
+	
+	# ? /Inference
+	ui_status_text.empty()
+	ui_progress_bar.empty()
+	
+	# ? Rehydrate the tiles into a full image.
+	complete_tiles_tensor = torch.from_numpy(numpy.stack(complete_tiles))
+	complete_image = image_tiler.rebuild_with_masks(complete_tiles_tensor)
+	
+	# ? Unpad the image, a simple crop.
+	if place_x > 0 or place_y > 0:
+		complete_image = complete_image[:, :, place_y:place_y + orig_h, place_x:place_x + orig_w]
+	
+	# ? Clamp and convert to PIL.
+	complete_image = complete_image.squeeze(0)
+	complete_image = complete_image.clamp(0, 1.0)
+	final_image_pil = TVTF.to_pil_image(complete_image)
+	
+	return final_image_pil
+
+def streamlit_to_pil_image(streamlit_file: UploadedFile):
+	image = Image.open(io.BytesIO(streamlit_file.read()))
+	return image
+
+def pil_to_buffered_png(image: Image.Image) -> io.BytesIO:
+	buffer = io.BytesIO()
+	image.save(buffer, format="PNG", compression=3)
+	buffer.seek(0)
+	return buffer
+
+# ! Image Inference
+if ui_image_input is not None and ui_image_input.name is not None:
+	bootstrap_model()
+	ui_status_text.empty()
+	ui_progress_bar.empty()
+	
+	onnx_session.get_modelmeta()
+	onnx_input_metadata = onnx_session.get_inputs()[0]
+	b, c, h, w = onnx_input_metadata.shape
+	
+	target_batch_size = b
+	# This is always square, if H and W are different for ONNX input you screwed up, so I don't want to hear it.
+	target_tile_size = h
+	
+	input_image = streamlit_to_pil_image(ui_image_input)
+	loaded_image_pt = image_tools.prepare_image_for_inference(input_image)
+	finished_image = evaluate_tiled(loaded_image_pt, tile_size=target_tile_size, batch_size=target_batch_size)
+	
+	with ui_col2:
+		ui_image_output.image(finished_image, use_container_width=True, caption="Output Image")
+		complete_file_name = f"{ui_image_input.name.rsplit('.', 1)[0]}_output.png"
+		
+		@streamlit.fragment
+		def download_button():
+			# ui_image_download.download_button("Download Image", image_to_bytesio(finished_image), complete_file_name, type="primary", on_click=lambda: setattr(streamlit.session_state, 'download_click', True))
+			streamlit.download_button("Download Image", pil_to_buffered_png(finished_image), complete_file_name, type="primary")
+		
+		download_button()
+	
+	with ui_section_compare:
+		image_comparison(img1=input_image, img2=finished_image, make_responsive=True, label1="Input Image", label2="Output Image", width=1024)

pycharm_runner.py

@@ -0,0 +1,9 @@
+# https://discuss.streamlit.io/t/cannot-debug-streamlit-in-pycharm-2023-3-3/61581/2
+
+try:
+	from streamlit.web import bootstrap
+except ImportError:
+	from streamlit import bootstrap
+
+real_script = 'app.py'
+bootstrap.run(real_script, is_hello=False, args=[], flag_options={})

requirements.txt

@@ -0,0 +1,9 @@
+blended_tiling==0.0.1.dev7
+huggingface-hub==0.27.1
+onnx==1.17.0
+onnxruntime==1.17.0
+pillow==11.0.0
+streamlit==1.41.1
+streamlit-image-comparison==0.0.4
+torch
+torchvision

tools/__init__.py

@@ -0,0 +1,3 @@
+
+if __name__ == '__main__':
+	print('__main__ is not supported in modules')

tools/image_tools.py

@@ -0,0 +1,43 @@
+import torch
+from PIL import Image
+from torchvision.transforms import functional as TVTF
+
+from .mafs import round_up_to_multiple
+
+def prepare_image_for_inference(image_pil: Image.Image) -> torch.Tensor:
+	if image_pil.mode != 'RGB':
+		image_pil = image_pil.convert('RGB')
+	
+	# Just being explicit, in case of environmental oddities or something.
+	image_pt = TVTF.to_tensor(image_pil)
+	image_pt = image_pt.to(device='cpu', dtype=torch.float32)
+	
+	return image_pt
+
+def pad_to_divisible(image_tensor: torch.Tensor, tile_size: int = 128):
+	c, h, w = image_tensor.shape
+	
+	# If the dims are already divisible by the tile size, we're good.
+	if h % tile_size == 0 and w % tile_size == 0:
+		return image_tensor, 0, 0
+	
+	expanded_w = round_up_to_multiple(w, tile_size)
+	expanded_h = round_up_to_multiple(h, tile_size)
+	l, t, r, b = 0, 0, 0, 0
+	
+	# Distribute the padding evenly on all sides.
+	if expanded_w > w:
+		diff = expanded_w - w
+		l = diff // 2
+		r = diff - l
+	if expanded_h > h:
+		diff = expanded_h - h
+		t = diff // 2
+		b = diff - t
+	
+	image_tensor = TVTF.pad(image_tensor, padding=[l, t, r, b], padding_mode='reflect')
+	
+	return image_tensor, l, t
+
+if __name__ == '__main__':
+	print('__main__ not supported in modules.')

tools/mafs.py

@@ -0,0 +1,5 @@
+def round_up_to_multiple(x: int, multiple: int) -> int:
+	return ((x + multiple - 1) // multiple) * multiple
+
+if __name__ == '__main__':
+	print('__main__ not supported in modules.')