updated app

app.py

@@ -1,34 +1,145 @@
-import streamlit as st
+import numpy as np
+import pandas as pd
 from PIL import Image
+from collections import defaultdict
+
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
 
 import torch
 from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
 from diffusers import StableDiffusionInpaintPipeline
 
+import matplotlib as mpl
+
+from model import segment_image, inpaint
+
+
+# define utils and helpers
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+def closest_number(n, m=8):
+    """ Obtains closest number to n that is divisble by m """
+    return int(n/m) * m
+
+
+def get_mask_from_rectangles(image, mask, height, width, drawing_mode='rect'):
+    # Create a canvas component
+    canvas_result = st_canvas(
+        fill_color="rgba(255, 165, 0, 0.3)",  
+        stroke_width=2,
+        stroke_color="#000000", 
+        background_image=image,
+        update_streamlit=True,
+        height=height,
+        width=width,
+        drawing_mode=drawing_mode,
+        point_display_radius=5,
+        key="canvas",
+    )
+
+    # get selections from mask
+    if canvas_result.json_data is not None:
+        objects = pd.json_normalize(canvas_result.json_data["objects"])
+        for col in objects.select_dtypes(include=["object"]).columns:
+            objects[col] = objects[col].astype("str")
+
+        if len(objects) > 0:
+            # st.dataframe(objects)
+
+            left_coords = objects.left.to_numpy()
+            top_coords = objects.top.to_numpy()
+            right_coords = left_coords + objects.width.to_numpy()
+            bottom_coords = top_coords + objects.height.to_numpy()
+
+            # add selections to mask
+            for (left, top, right, bottom) in zip(left_coords, top_coords, right_coords, bottom_coords):
+                cropped = image.crop((left, top, right, bottom))
+                st.image(cropped)
+                mask[top:bottom, left:right] = 255
+
+            st.header("Mask Created!")
+            st.image(mask)
+
+    return mask
+
+
+def get_mask(image, edit_method, height, width):
+    mask = np.zeros((height, width), dtype=np.uint8)
+
+    if edit_method == "AutoSegment Area":
+
+        # get displayable segmented image
+        seg_prediction, segment_labels = segment_image(image)
+        seg = seg_prediction['segmentation'].cpu().numpy()
+        viridis = mpl.colormaps.get_cmap('viridis').resampled(np.max(seg))
+        seg_image = Image.fromarray(np.uint8(viridis(seg)*255))
+
+        # display image
+        st.image(seg_image)
+
+        # prompt user to select valid labels to edit
+        seg_selections = st.multiselect("Choose segments", zip(segment_labels.keys(), segment_labels.values()))
+        if seg_selections:
+            tgts = []
+            for s in seg_selections:
+                tgts.append(s[0])
+
+            mask = Image.fromarray(np.array([(seg == t) for t in tgts]).sum(axis=0).astype(np.uint8)*255)
+            st.header("Mask Created!")
+            st.image(mask)
+
+    elif edit_method == "Draw Custom Area":
+        mask = get_mask_from_rectangles(image, mask, height, width)
+
+
+    return mask
+
+
+
+if __name__ == '__main__':
+    
+    st.title("Stable Edit - Edit your photos with Stable Diffusion!")
+
+    # upload image
+    filename = st.file_uploader("upload an image")
+    # filename = r"C:\Users\itber\Downloads\Fjord_Cycling.jpg"
+
+    sf = st.text_input("Please enter resizing scale factor to downsize image (default = 2)", value="2")
+    try: 
+        sf = int(sf)
+    except:
+        sf.write("Error with input scale factor, setting to default value of 2, please re-enter above to change it")
+        sf = 2
+
+    if filename:
+        image = Image.open(filename)
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        width, height = image.size
+        width, height = closest_number(width/sf), closest_number(height/sf)
+        image = image.resize((width, height))
 
-seg_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
-seg_model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+        st.image(image)
+        # st.write(f"{width} {height}")
 
+        # Select an editing method
+        edit_method = st.selectbox("Select Edit Method", ("AutoSegment Area", "Draw Custom Area"))
 
-# get Stable Diffusion model for image inpainting
-pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    "runwayml/stable-diffusion-inpainting",
-    torch_dtype=torch.float16,
-).to(device)
+        if edit_method:
+            mask = get_mask(image, edit_method, height, width)
 
+            # get inpainted images
+            prompt = st.text_input("Please enter prompt for image inpainting", value="")
 
-st.title("Stable Edit - Edit your photos with Stable Diffusion!")
+            st.write("Inpainting Images, patience is a virtue :)")
 
-# upload image
-filename = st.file_uploader("upload an image")
-image = Image.open(filename)
-st.image(image)
+            images = inpaint(image, mask, width, height, prompt=prompt, seed=0, guidance_scale=17.5, num_samples=3)
 
-# Select Area to edit
-selection = st.selectbox("Select Area(s) to edit", ("AutoSegment Area", "Draw Custom Area", "Suprise Me"))
+            # display all images
+            st.write("Original Image")
+            st.image(image)
+            for i, img in enumerate(images, 1):
+                st.write(f"result: {i}")
+                st.image(img)
 
-# TEMP - DEMO stuff
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+    

model.py

@@ -0,0 +1,72 @@
+import torch
+from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
+from diffusers import StableDiffusionInpaintPipeline
+
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+# Image segmentation
+seg_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+seg_model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+
+def segment_image(image):
+    inputs = seg_processor(image, return_tensors="pt")
+    
+    with torch.no_grad():
+        seg_outputs = seg_model(**inputs)
+
+    # get prediction dict
+    seg_prediction = seg_processor.post_process_panoptic_segmentation(seg_outputs, target_sizes=[image.size[::-1]])[0]
+
+    # get segment labels dict
+    segment_labels = {}
+    for segment in seg_prediction['segments_info']:
+        segment_id = segment['id']
+        segment_label_id = segment['label_id']
+        segment_label = seg_model.config.id2label[segment_label_id]
+
+        segment_labels.update({segment_id : segment_label})
+
+    return seg_prediction, segment_labels
+
+# Image inpainting
+# get Stable Diffusion model for image inpainting
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    torch_dtype=torch.float16,
+).to(device)
+
+
+def inpaint(image, mask, W, H, prompt="", seed=0, guidance_scale=17.5, num_samples=3):
+    """ Uses Stable Diffusion model to inpaint image 
+        Inputs: 
+            image - input image (PIL or torch tensor)
+            mask - mask for inpainting same size as image (PIL or troch tensor)
+            W - size of image
+            H - size of mask
+            prompt - prompt for inpainting
+            seed - random seed 
+        Outputs:
+            images - output images
+        """
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    images = pipe(
+        prompt=prompt,
+        image=image,
+        mask_image=mask, # ensure mask is same type as image
+        height=H,
+        width=W,
+        guidance_scale=guidance_scale,
+        generator=generator,
+        num_images_per_prompt=num_samples,
+    ).images
+    
+    return images
+
+    
+
+    
+
+
+