inpainting

.idea/$CACHE_FILE$

@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectInspectionProfilesVisibleTreeState">
+    <entry key="Project Default">
+      <profile-state>
+        <expanded-state>
+          <State>
+            <id />
+          </State>
+          <State>
+            <id>General</id>
+          </State>
+        </expanded-state>
+        <selected-state>
+          <State>
+            <id>AngularJS</id>
+          </State>
+        </selected-state>
+      </profile-state>
+    </entry>
+  </component>
+</project>

.idea/dictionaries

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectDictionaryState">
+    <dictionary name="running_you" />
+  </component>
+</project>

.idea/encodings.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="Encoding" defaultCharsetForPropertiesFiles="UTF-8">
+    <file url="PROJECT" charset="UTF-8" />
+  </component>
+</project>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,5 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="PROJECT_PROFILE" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="JavaScriptSettings">
+    <option name="languageLevel" value="ES6" />
+  </component>
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/new_inpainting.iml" filepath="$PROJECT_DIR$/.idea/new_inpainting.iml" />
+    </modules>
+  </component>
+</project>

.idea/new_inpainting.iml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="jdk" jdkName="Python 3.7 (modelscope)" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/workspace.xml

@@ -0,0 +1,18 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectId" id="2Id7cZROeBLwSmWyYcqsSXcYO5h" />
+  <component name="PropertiesComponent">
+    <property name="settings.editor.selected.configurable" value="editor.preferences.import" />
+  </component>
+  <component name="PyConsoleOptionsProvider">
+    <option name="myPythonConsoleState">
+      <console-settings sdk-home="$USER_HOME$/anaconda3/bin/python" working-directory="$PROJECT_DIR$/..">
+        <option name="mySdkHome" value="$USER_HOME$/anaconda3/bin/python" />
+        <option name="myWorkingDirectory" value="$PROJECT_DIR$/.." />
+      </console-settings>
+    </option>
+  </component>
+  <component name="VcsContentAnnotationSettings">
+    <option name="myLimit" value="2678400000" />
+  </component>
+</project>

app.py

@@ -1,8 +1,94 @@
+import imageio
+from PIL import Image
+import math
+import streamlit as st
+import numpy as np
+import torch
+import PIL
+import cv2
+import mediapipe as mp
 import gradio as gr
+from diffusers import StableDiffusionInpaintPipeline
 
-def greet(name):
-    return "Hello " + name + "!!"
+YOUR_TOKEN = st.secrets['USER_TOKEN']
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model_path = "runwayml/stable-diffusion-inpainting"
 
+if torch.cuda.is_available():
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(model_path, revision="fp16", torch_dtype=torch.float16,
+                                                          use_auth_token=YOUR_TOKEN).to(device)
+else:
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(model_path, use_auth_token=YOUR_TOKEN).to(device)
+
+
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows * cols
+
+    w, h = imgs[0].size
+    grid = PIL.Image.new('RGB', size=(cols * w, rows * h))
+    grid_w, grid_h = grid.size
+
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+
+
+def mediapipe_segmentation(image_file, mask_file):
+    mp_drawing = mp.solutions.drawing_utils
+    mp_selfie_segmentation = mp.solutions.selfie_segmentation
+
+    # For static images:
+    BG_COLOR = (0, 0, 0)  # gray
+    MASK_COLOR = (255, 255, 255)  # white
+    with mp_selfie_segmentation.SelfieSegmentation(model_selection=0) as selfie_segmentation:
+        image = cv2.imread(image_file)
+        image_height, image_width, _ = image.shape
+        # Convert the BGR image to RGB before processing.
+        results = selfie_segmentation.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+        # blurred_image = cv2.GaussianBlur(image,(55,55),0)
+        # condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
+        # output_image = np.where(condition, image, blurred_image)
+
+        # Draw selfie segmentation on the background image.
+        # To improve segmentation around boundaries, consider applying a joint
+        # bilateral filter to "results.segmentation_mask" with "image".
+        condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
+        # Generate solid color images for showing the output selfie segmentation mask.
+        fg_image = np.zeros(image.shape, dtype=np.uint8)
+        fg_image[:] = MASK_COLOR
+        bg_image = np.zeros(image.shape, dtype=np.uint8)
+        bg_image[:] = BG_COLOR
+        output_image = np.where(condition, fg_image, bg_image)
+        cv2.imwrite(mask_file, output_image)
+
+
+def image_inpainting(prompt, image_path, mask_image_path, num_samples=4):
+    image = PIL.Image.open(image_path).convert("RGB").resize((512, 512))
+    mask_image = PIL.Image.open(mask_image_path).convert("RGB").resize((512, 512))
+
+    guidance_scale = 7.5
+    generator = torch.Generator(device=device).manual_seed(0)  # change the seed to get different results
+
+    images = pipe(prompt=prompt, image=image, mask_image=mask_image, guidance_scale=guidance_scale, generator=generator,
+                  num_images_per_prompt=num_samples).images
+
+    # insert initial image in the list so we can compare side by side
+    # images.insert(0, image)
+    return image_grid(images, 2, math.ceil(num_samples/2))
+
+
+def predict1(dict, prompt):
+    dict['image'].save('image.png')
+    dict['mask'].save('mask.png')
+    mediapipe_segmentation('image.png', 'm_mask.png')
+
+    image = image_inpainting(prompt, image_path='image.png', mask_image_path='m_mask.png')
+    return image
+
+
+title = "Person Matting & Stable Diffusion In-Painting"
+description = "Inpainting Stable Diffusion <br/><b>mediapipe + Stable Diffusion<b/><br/>"
+gr.Interface(predict1, inputs=[gr.Image(source='upload', tool='sketch', type='pil'), gr.Textbox(label='prompt')],
+             outputs='image', title=title, description=description).launch(max_threads=True)

requirements.txt

@@ -0,0 +1,8 @@
+diffusers>=0.6.0
+transformers
+mediapipe
+scipy
+ftfy
+imageio
+torch
+streamlit