Spaces:

jhj0517
/

Segment-Anything-Layer-Divider

Build error

App Files Files Community

jhj0517 commited on Apr 9, 2023

Commit

11d7b39

•

1 Parent(s): 4af6ba2

initial commit

Browse files

Files changed (12) hide show

.gitignore +3 -0
app.py +49 -0
models/model file will be saved here.txt +0 -0
modules/__init__.py +0 -0
modules/html_constants.py +84 -0
modules/mask_utils.py +99 -0
modules/model_downloader.py +12 -0
modules/sam.py +65 -0
modules/ui_utils.py +8 -0
outputs/psd/psd file will be saved here.txt +0 -0
requirements.txt +12 -0
screenshot.png +0 -0

.gitignore ADDED Viewed

	@@ -0,0 +1,3 @@

+models/
+modules/__pycache__/
+outputs/

app.py ADDED Viewed

	@@ -0,0 +1,49 @@

+import gradio as gr
+from modules import sam
+from modules.ui_utils import *
+from modules.html_constants import *
+class App:
+    def __init__(self):
+        #download_sam_model_url()
+        self.app = gr.Blocks(css=CSS)
+        self.sam = sam.SamInference()
+    def launch(self):
+        with self.app:
+            with gr.Row():
+                gr.Markdown(MARKDOWN_NOTE, elem_id="md_pgroject")
+            with gr.Row().style(equal_height=True):  # bug https://github.com/gradio-app/gradio/issues/3202
+                with gr.Column(scale=5):
+                    img_input = gr.Image(label="Input image here")
+                with gr.Column(scale=5):
+                    # Tuable Params
+                    nb_points_per_side = gr.Number(label="points_per_side", value=32)
+                    sld_pred_iou_thresh = gr.Slider(label="pred_iou_thresh", value=0.88, minimum=0, maximum=1)
+                    sld_stability_score_thresh = gr.Slider(label="stability_score_thresh", value=0.95, minimum=0,
+                                                           maximum=1)
+                    nb_crop_n_layers = gr.Number(label="crop_n_layers", value=0)
+                    nb_crop_n_points_downscale_factor = gr.Number(label="crop_n_points_downscale_factor", value=1)
+                    nb_min_mask_region_area = gr.Number(label="min_mask_region_area", value=0)
+                    html_param_explain = gr.HTML(PARAMS_EXPLANATION, elem_id="html_param_explain")
+            with gr.Row():
+                btn_generate = gr.Button("GENERATE", variant="primary")
+            with gr.Row():
+                gallery_output = gr.Gallery(label="Output will be shown here", show_label=True).style(grid=5,
+                                                                                                      height="auto")
+                btn_open_folder = gr.Button("📁\n(PSD)").style(full_width=False)
+            params = [nb_points_per_side, sld_pred_iou_thresh, sld_stability_score_thresh, nb_crop_n_layers,
+                      nb_crop_n_points_downscale_factor, nb_min_mask_region_area]
+            btn_generate.click(fn=self.sam.generate_mask_app, inputs=[img_input] + params, outputs=gallery_output)
+            btn_open_folder.click(fn=lambda: open_folder("outputs\psd"), inputs=None, outputs=None)
+        self.app.queue(api_open=False).launch()
+if __name__ == "__main__":
+    app = App()
+    app.launch()

models/model file will be saved here.txt ADDED Viewed

File without changes

modules/__init__.py ADDED Viewed

File without changes

modules/html_constants.py ADDED Viewed

	@@ -0,0 +1,84 @@

+CSS = """
+#md_project a {
+  color: black;
+  text-decoration: none;
+}
+#md_project a:hover {
+  text-decoration: underline;
+}
+"""
+PROJECT_NAME = """
+# [Layer-Divider-WebUI](https://github.com/jhj0517/Layer-Divider-WebUI)
+"""
+MARKDOWN_NOTE = """
+## This space only support CPU because it's free huggingface space.
+## If you want to run CUDA version , check this [repository](https://github.com/jhj0517/Layer-Divider-WebUI)
+"""
+PARAMS_EXPLANATION = """
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <style>
+    table {
+      border-collapse: collapse;
+      width: 100%;
+    }
+    th, td {
+      border: 1px solid #dddddd;
+      text-align: left;
+      padding: 8px;
+    }
+    th {
+      background-color: #f2f2f2;
+    }
+  </style>
+</head>
+<body>
+<details>
+  <summary>Explanation of Each Parameter</summary>
+  <table>
+    <thead>
+      <tr>
+        <th>Parameter</th>
+        <th>Description</th>
+      </tr>
+    </thead>
+    <tbody>
+      <tr>
+        <td>points_per_side</td>
+        <td>The number of points to be sampled along one side of the image. The total number of points is points_per_side**2. If None, 'point_grids' must provide explicit point sampling.</td>
+      </tr>
+      <tr>
+        <td>pred_iou_thresh</td>
+        <td>A filtering threshold in [0,1], using the model's predicted mask quality.</td>
+      </tr>
+      <tr>
+        <td>stability_score_thresh</td>
+        <td>A filtering threshold in [0,1], using the stability of the mask under changes to the cutoff used to binarize the model's mask predictions.</td>
+      </tr>
+      <tr>
+        <td>crops_n_layers</td>
+        <td>If >0, mask prediction will be run again on crops of the image. Sets the number of layers to run, where each layer has 2**i_layer number of image crops.</td>
+      </tr>
+      <tr>
+        <td>crop_n_points_downscale_factor</td>
+        <td>The number of points-per-side sampled in layer n is scaled down by crop_n_points_downscale_factor**n.</td>
+      </tr>
+      <tr>
+        <td>min_mask_region_area</td>
+        <td>If >0, postprocessing will be applied to remove disconnected regions and holes in masks with area smaller than min_mask_region_area. Requires opencv.</td>
+      </tr>
+    </tbody>
+  </table>
+</details>
+</body>
+</html>
+"""

modules/mask_utils.py ADDED Viewed

	@@ -0,0 +1,99 @@

+import cv2
+import numpy as np
+from pycocotools import mask as coco_mask
+from pytoshop import layers
+import pytoshop
+from pytoshop.enums import BlendMode
+from datetime import datetime
+def generate_random_color():
+    return np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256)
+def create_base_layer(image):
+    rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
+    return [rgba_image]
+def create_mask_layers(image, masks):
+    layer_list = []
+    for result in masks:
+        rle = result['segmentation']
+        mask = coco_mask.decode(rle).astype(np.uint8)
+        rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
+        rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)
+        layer_list.append(rgba_image)
+    return layer_list
+def create_mask_gallery(image, masks):
+    mask_array_list = []
+    label_list = []
+    for index, result in enumerate(masks):
+        rle = result['segmentation']
+        mask = coco_mask.decode(rle).astype(np.uint8)
+        rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
+        rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)
+        mask_array_list.append(rgba_image)
+        label_list.append(f'Part {index}')
+    return [[img, label] for img, label in zip(mask_array_list, label_list)]
+def create_mask_combined_images(image, masks):
+    final_result = np.zeros_like(image)
+    for result in masks:
+        rle = result['segmentation']
+        mask = coco_mask.decode(rle).astype(np.uint8)
+        color = generate_random_color()
+        colored_mask = np.zeros_like(image)
+        colored_mask[mask == 1] = color
+        final_result = cv2.addWeighted(final_result, 1, colored_mask, 0.5, 0)
+    combined_image = cv2.addWeighted(image, 1, final_result, 0.5, 0)
+    return [combined_image, "masked"]
+def insert_psd_layer(psd, image_data, layer_name, blending_mode):
+    channel_data = [layers.ChannelImageData(image=image_data[:, :, i], compression=1) for i in range(4)]
+    layer_record = layers.LayerRecord(
+        channels={-1: channel_data[3], 0: channel_data[0], 1: channel_data[1], 2: channel_data[2]},
+        top=0, bottom=image_data.shape[0], left=0, right=image_data.shape[1],
+        blend_mode=blending_mode,
+        name=layer_name,
+        opacity=255,
+    )
+    psd.layer_and_mask_info.layer_info.layer_records.append(layer_record)
+    return psd
+def save_psd(input_image_data, layer_data, layer_names, blending_modes):
+    psd_file = pytoshop.core.PsdFile(num_channels=3, height=input_image_data.shape[0], width=input_image_data.shape[1])
+    for index, layer in enumerate(layer_data):
+        psd_file = insert_psd_layer(psd_file, layer, layer_names[index], blending_modes[index])
+    timestamp = datetime.now().strftime("%m%d%H%M%S")
+    with open(f"outputs/psd/result-{timestamp}.psd", 'wb') as output_file:
+        psd_file.write(output_file)
+def save_psd_with_masks(image, masks):
+    original_layer = create_base_layer(image)
+    mask_layers = create_mask_layers(image, masks)
+    names = [f'Part {i}' for i in range(len(mask_layers))]
+    modes = [BlendMode.normal] * (len(mask_layers)+1)
+    save_psd(image, original_layer+mask_layers, ['Original_Image']+names, modes)

modules/model_downloader.py ADDED Viewed

	@@ -0,0 +1,12 @@

+import torch
+AVAILABLE_MODELS = {
+    "ViT-H SAM model": ["sam_vit_h_4b8939.pth", "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"],
+    "ViT-L SAM model": ["sam_vit_l_0b3195.pth", "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth"],
+    "ViT-B SAM model": ["sam_vit_b_01ec64.pth", "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"],
+}
+def download_sam_model_url():
+    torch.hub.download_url_to_file(AVAILABLE_MODELS["ViT-H SAM model"][1],
+                                   f'models/{AVAILABLE_MODELS["ViT-H SAM model"][0]}')

modules/sam.py ADDED Viewed

	@@ -0,0 +1,65 @@

+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
+import os
+from modules.mask_utils import *
+from modules.model_downloader import *
+class SamInference:
+    def __init__(self):
+        self.model = None
+        self.model_path = f"models/sam_vit_h_4b8939.pth"
+        self.device = "cuda"
+        self.mask_generator = None
+        # Tuable Parameters , All default values
+        self.tunable_params = {
+            'points_per_side': 32,
+            'pred_iou_thresh': 0.88,
+            'stability_score_thresh': 0.95,
+            'crop_n_layers': 0,
+            'crop_n_points_downscale_factor': 1,
+            'min_mask_region_area': 0
+        }
+    def set_mask_generator(self):
+        print("applying configs to model..")
+        if not os.path.exists(self.model_path):
+            print("No needed SAM model detected. downloading VIT H SAM model....")
+            download_sam_model_url()
+        self.model = sam_model_registry["default"](checkpoint=self.model_path)
+        self.model.to(device=self.device)
+        self.mask_generator = SamAutomaticMaskGenerator(
+            self.model,
+            points_per_side=self.tunable_params['points_per_side'],
+            pred_iou_thresh=self.tunable_params['pred_iou_thresh'],
+            stability_score_thresh=self.tunable_params['stability_score_thresh'],
+            crop_n_layers=self.tunable_params['crop_n_layers'],
+            crop_n_points_downscale_factor=self.tunable_params['crop_n_points_downscale_factor'],
+            min_mask_region_area=self.tunable_params['min_mask_region_area'],
+            output_mode="coco_rle",
+        )
+    def generate_mask(self, image):
+        return [self.mask_generator.generate(image)]
+    def generate_mask_app(self, image, *params):
+        tunable_params = {
+            'points_per_side': int(params[0]),
+            'pred_iou_thresh': float(params[1]),
+            'stability_score_thresh': float(params[2]),
+            'crop_n_layers': int(params[3]),
+            'crop_n_points_downscale_factor': int(params[4]),
+            'min_mask_region_area': int(params[5]),
+        }
+        if self.model is None or self.mask_generator is None or self.tunable_params != tunable_params:
+            self.tunable_params = tunable_params
+            self.set_mask_generator()
+        masks = self.mask_generator.generate(image)
+        save_psd_with_masks(image, masks)
+        combined_image = create_mask_combined_images(image, masks)
+        gallery = create_mask_gallery(image, masks)
+        return [combined_image] + gallery

modules/ui_utils.py ADDED Viewed

	@@ -0,0 +1,8 @@

+import os
+def open_folder(folder_path):
+    if os.path.exists(folder_path):
+        os.system(f"start {folder_path}")
+    else:
+        print(f"The folder {folder_path} does not exist.")

outputs/psd/psd file will be saved here.txt ADDED Viewed

File without changes

requirements.txt ADDED Viewed

	@@ -0,0 +1,12 @@

+--extra-index-url https://download.pytorch.org/whl/cu117
+torch
+--extra-index-url https://download.pytorch.org/whl/cu117
+torchvision
+git+https://github.com/facebookresearch/segment-anything.git
+opencv-python
+pycocotools
+matplotlib
+onnxruntime
+onnx
+gradio
+pytoshop==1.2.0

screenshot.png ADDED Viewed