first commit

app.py

@@ -0,0 +1,146 @@
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
+from PIL import Image
+import numpy as np
+import random
+import vipainting
+import time
+import threading
+from queue import Queue
+import os
+
+image_queue = Queue()
+sampling_queue = Queue()
+
+
+st.title("Mask Your Own Inpaint")
+
+@st.cache_data
+def load_images():
+    data = np.load("data/sflckr_all_images.npz")
+    images = data["images"]  
+    return images
+
+if "random_idx" not in st.session_state:
+    st.session_state.random_idx = None
+
+images = load_images()
+if st.button("Random Pick"):
+    st.session_state.random_idx = random.randint(0, images.shape[0] - 1)
+
+def make_square(img, target_size=300):
+    size = max(img.size)  
+    background = Image.new("RGB", (size, size), (255, 255, 255))
+    offset = ((size - img.size[0]) // 2, (size - img.size[1]) // 2)
+    background.paste(img, offset)
+    return background.resize((target_size, target_size))
+
+def run_inpainting(random_idx, mask_array, image_queue, sampling_queue):
+    vipainting.vipaint(random_idx, mask_array, image_queue, sampling_queue)
+
+
+if st.session_state.random_idx is not None:
+    img_array = images[st.session_state.random_idx]
+
+    img_pil = Image.fromarray(img_array)
+    img_pil = make_square(img_pil, target_size=300)  
+
+
+    col1, col2 = st.columns(2)
+    with col1:
+        st.write("Draw your mask on the image below:")
+        canvas_result = st_canvas(
+            fill_color="rgba(255, 0, 0, 0.3)",  
+            stroke_width=50,
+            stroke_color="black",
+            background_image=img_pil,
+            update_streamlit=True,
+            width=300,
+            height=300,
+            drawing_mode="freedraw",
+            key="canvas"
+        )
+
+
+    if canvas_result.image_data is not None:
+        mask = canvas_result.image_data[:, :, 3]  
+        binary_mask = (mask > 128).astype(np.uint8) * 255
+
+        with col2:
+            st.write("Masked Image")
+            st.image(binary_mask, caption="Binary Mask", width=300)  
+
+        mask_image = Image.fromarray(binary_mask).resize((512, 512), Image.ANTIALIAS)
+        mask_array = 255 - np.array(mask_image)
+        mask_array = np.expand_dims(mask_array, axis=(0, 1)) 
+
+        if st.button("inpaint"):
+            st.write("Please wait...")
+            inpaint_thread = threading.Thread(target=run_inpainting, args=(st.session_state.random_idx, mask_array, image_queue, sampling_queue))
+            inpaint_thread.start()
+
+            img_left, img_right = st.columns(2) 
+            img_left_placeholder = img_left.empty()
+            img_right_placeholder = img_right.empty()
+            with img_left:
+                img_left_placeholder.image(img_pil, caption=f"True Image", width=300)
+            seg_image_path = f"results/{st.session_state.random_idx}/input.png"
+
+            while True:
+                if os.path.exists(seg_image_path):
+                    with img_right:
+                        img_right_image = Image.open(seg_image_path)
+                        img_right_placeholder.image(img_right_image, caption="Segmentation Map", width=300)
+                    break  
+                time.sleep(0.5)  
+
+
+            # Set up progress tracking
+            expected_updates = 100
+            progress_bar = st.progress(0)
+            st.write("Fitting in progress")
+            displayed_images = 0
+
+            col_left, col_right = st.columns(2)
+            left_placeholder = col_left.empty()
+            right_placeholder = col_right.empty()
+
+            
+            while displayed_images < expected_updates:
+                if not image_queue.empty():
+                    img = image_queue.get()  # Get the next image from the queue
+
+                    if displayed_images % 2 == 0:
+                        left_placeholder.image(img, caption=f"Progress Update {displayed_images + 1}", width=300)
+                    else:
+                        right_placeholder.image(img, caption=f"Progress Update {displayed_images + 1}", width=300)
+
+                    # Update progress bar
+                    displayed_images += 1
+                    progress_bar.progress(displayed_images / expected_updates)
+
+                time.sleep(0.3)
+
+            expected_updates = 10
+            s_progress_bar = st.progress(0)
+            displayed_images = 0
+            st.write("Sampling in progress")
+            sample_left, sample_right = st.columns(2)
+            sleft_placeholder = sample_left.empty()
+            sright_placeholder = sample_right.empty()
+            while displayed_images < expected_updates:
+                if not sampling_queue.empty():
+                    img = sampling_queue.get()  
+
+                    if displayed_images % 2 == 0:
+                        sleft_placeholder.image(img, caption=f"Sampling Update {displayed_images + 1}", width=300)
+                    else:
+                        sright_placeholder.image(img, caption=f"Sampling Update {displayed_images + 1}", width=300)
+
+                    displayed_images += 1
+                    s_progress_bar.progress(displayed_images / expected_updates)
+
+                time.sleep(0.3)
+
+            inpaint_thread.join()
+            st.success("Inpainting completed!")

requirements.txt

@@ -0,0 +1,115 @@
+streamlit
+streamlit-drawable-canvas
+pillow
+numpy
+# Core Packages
+numpy==1.24.4
+pillow==9.5.0
+torch==2.4.1
+torchvision==0.8.1
+opencv-python==4.10.0.84
+opencv-python-headless==4.10.0.84
+tqdm==4.66.5
+scipy==1.10.1
+pandas==2.0.3
+matplotlib==3.7.5
+streamlit==1.39.0
+streamlit-drawable-canvas==0.9.3
+
+# PyPI Packages
+absl-py==2.1.0
+aiohttp==3.10.10
+aiohappyeyeballs==2.4.3
+aiosignal==1.3.1
+albumentations==1.4.18
+altair==5.4.1
+async-timeout==4.0.3
+attrs==24.2.0
+blinker==1.8.2
+cachetools==5.5.0
+charset-normalizer==3.4.0
+click==8.1.7
+contourpy==1.1.1
+diffusers==0.31.0
+docker-pycreds==0.4.0
+einops==0.8.0
+filelock==3.16.1
+fonttools==4.54.1
+fsspec==2024.10.0
+gitdb==4.0.11
+gitpython==3.1.43
+google-auth==2.35.0
+google-auth-oauthlib==1.0.0
+grpcio==1.67.0
+huggingface-hub==0.26.1
+idna==3.10
+imageio==2.35.1
+importlib-metadata==8.5.0
+importlib-resources==6.4.5
+invisible-watermark==0.2.0
+jinja2==3.1.4
+jsonschema==4.23.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.7
+kornia==0.6.4
+markdown==3.7
+markdown-it-py==3.0.0
+matplotlib==3.7.5
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.1.0
+networkx==3.1
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.6.77
+nvidia-nvtx-cu12==12.1.105
+oauthlib==3.2.2
+omegaconf==2.3.0
+packaging==24.1
+pkgutil-resolve-name==1.3.10
+protobuf==3.20.1
+psutil==6.1.0
+pyarrow==17.0.0
+pydeck==0.9.1
+pydeprecate==0.3.2
+pygments==2.18.0
+pyparsing==3.1.4
+python-dateutil==2.9.0.post0
+pytorch-lightning==1.6.5
+pyyaml==6.0.2
+referencing==0.35.1
+regex==2024.9.11
+requests==2.32.3
+requests-oauthlib==2.0.0
+rich==13.9.3
+rsa==4.9
+safetensors==0.4.5
+scikit-image==0.21.0
+sentry-sdk==2.17.0
+setproctitle==1.3.3
+smmap==5.0.1
+sympy==1.13.3
+taming-transformers-rom1504==0.0.6
+tenacity==9.0.0
+tensorboard==2.14.0
+tensorboard-data-server==0.7.2
+tifffile==2023.7.10
+tokenizers==0.12.1
+toml==0.10.2
+torchmetrics==0.6.0
+transformers==4.19.2
+triton==3.0.0
+urllib3==2.2.3
+wandb==0.18.5
+watchdog==4.0.2
+werkzeug==3.0.6
+yarl==1.15.2
+zipp==3.20.2