Check dependency

app.py

@@ -1,49 +1,25 @@
 import gradio as gr
 import numpy as np
 
+from sam_segment import predict_masks_with_sam
+from lama_inpaint import inpaint_img_with_lama
+from utils import load_img_to_array, save_array_to_img, dilate_mask, \
+    show_mask, show_points
+
+
 with gr.Blocks() as demo:
-    tolerance = gr.Slider(label="Tolerance",
-                          info="How different colors can be in a segment.",
-                          minimum=0, maximum=256 * 3, value=50)
     with gr.Row():
         input_img = gr.Image(label="Input")
         output_img = gr.Image(label="Selected Segment")
 
+    with gr.Row():
+        h = gr.Number()
+        w = gr.Number()
 
-    def get_select_coords(img, tolerance, evt: gr.SelectData):
-        visited_pixels = set()
-        pixels_in_queue = set()
-        pixels_in_segment = set()
-        start_pixel = img[evt.index[1], evt.index[0]]
-        pixels_in_queue.add((evt.index[1], evt.index[0]))
-        while len(pixels_in_queue) > 0:
-            pixel = pixels_in_queue.pop()
-            visited_pixels.add(pixel)
-            neighbors = []
-            if pixel[0] > 0:
-                neighbors.append((pixel[0] - 1, pixel[1]))
-            if pixel[0] < img.shape[0] - 1:
-                neighbors.append((pixel[0] + 1, pixel[1]))
-            if pixel[1] > 0:
-                neighbors.append((pixel[0], pixel[1] - 1))
-            if pixel[1] < img.shape[1] - 1:
-                neighbors.append((pixel[0], pixel[1] + 1))
-            for neighbor in neighbors:
-                if neighbor in visited_pixels:
-                    continue
-                neighbor_pixel = img[neighbor[0], neighbor[1]]
-                if np.abs(neighbor_pixel - start_pixel).sum() < tolerance:
-                    pixels_in_queue.add(neighbor)
-                    pixels_in_segment.add(neighbor)
-
-        out = img.copy() * 0.2
-        out = out.astype(np.uint8)
-        for pixel in pixels_in_segment:
-            out[pixel[0], pixel[1]] = img[pixel[0], pixel[1]]
-        return out
-
+    def get_select_coords(img, evt: gr.SelectData):
+        return evt.index[1], evt.index[0]
 
-    input_img.select(get_select_coords, [input_img, tolerance], output_img)
+    input_img.select(get_select_coords, [input_img,], [h, w])
 
 if __name__ == "__main__":
     demo.launch()

lama_inpaint.py

@@ -16,6 +16,7 @@ os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
 os.environ['NUMEXPR_NUM_THREADS'] = '1'
 
 sys.path.insert(0, str(Path(__file__).resolve().parent / "third_party" / "lama"))
+
 from saicinpainting.evaluation.utils import move_to_device
 from saicinpainting.training.trainers import load_checkpoint
 from saicinpainting.evaluation.data import pad_tensor_to_modulo

requirements.txt

@@ -0,0 +1,23 @@
+torch
+torchvision
+opencv-python
+matplotlib
+tensorflow
+pyyaml
+tqdm
+numpy
+easydict
+scikit-image
+scikit-learn
+opencv-python
+joblib
+matplotlib
+pandas
+albumentations==0.5.2
+hydra-core
+pytorch-lightning
+tabulate
+kornia==0.5.0
+webdataset
+packaging
+wldhx.yadisk-direct

sam_segment.py

@@ -6,6 +6,8 @@ from matplotlib import pyplot as plt
 from typing import Any, Dict, List
 import torch
 
+sys.path.insert(0, str(Path(__file__).resolve().parent / "third_party" / "segment-anything"))
+
 from segment_anything import SamPredictor, sam_model_registry
 from utils import load_img_to_array, save_array_to_img, dilate_mask, \
     show_mask, show_points