antoyo123 commited on
Commit
032ff24
·
1 Parent(s): 063f21b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +58 -0
app.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
2
+ import gradio as gr
3
+ from PIL import Image
4
+ import matplotlib.pyplot as plt
5
+ import torch
6
+ import cv2
7
+
8
+
9
+ import os
10
+ os.system("wget https://huggingface.co/akhaliq/lama/resolve/main/best.ckpt")
11
+ import paddlehub as hub
12
+ import gradio as gr
13
+ import torch
14
+ from PIL import Image, ImageOps
15
+ import numpy as np
16
+ import imageio
17
+ os.mkdir("data")
18
+ os.rename("best.ckpt", "models/best.ckpt")
19
+ os.mkdir("dataout")
20
+
21
+ # Load CLIPSeg model
22
+ processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
23
+ clipseg_model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
24
+
25
+ # Load LAMA model
26
+ lama_model = hub.Module(name='U2Net')
27
+
28
+ def process_image(image, prompt):
29
+ # Generate mask with CLIPSeg
30
+ inputs = processor(text=prompt, images=image, padding="max_length", return_tensors="pt")
31
+ with torch.no_grad():
32
+ outputs = clipseg_model(**inputs)
33
+ preds = outputs.logits
34
+ plt.imsave("mask.png", torch.sigmoid(preds))
35
+ mask_image = Image.open("mask.png").convert("RGB")
36
+
37
+ # Convert image to BGR format
38
+ image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
39
+ # Convert mask to grayscale format
40
+ mask_image = cv2.cvtColor(np.array(mask_image), cv2.COLOR_RGB2GRAY)
41
+
42
+ # Perform inpainting with LAMA
43
+ input_dict = {"image": image, "mask": mask_image}
44
+ inpainted_image = lama_model.inference(data=input_dict)["data"][0]
45
+
46
+ inpainted_image = cv2.cvtColor(inpainted_image, cv2.COLOR_BGR2RGB)
47
+ inpainted_image = Image.fromarray(inpainted_image)
48
+
49
+ return mask_image, inpainted_image
50
+
51
+ interface = gr.Interface(fn=process_image,
52
+ inputs=[gr.Image(type="pil"), gr.Textbox(label="Please describe what you want to identify")],
53
+ outputs=[gr.Image(type="pil"), gr.Image(type="pil")],
54
+ title="Interactive demo: zero-shot image segmentation with CLIPSeg and inpainting with LAMA",
55
+ description="Demo for using CLIPSeg and LAMA to perform zero- and one-shot image segmentation and inpainting. To use it, simply upload an image and add a text to mask (identify in the image), or use one of the examples below and click 'submit'. Results will show up in a few seconds.",
56
+ examples=[["example_image.png", "wood"]])
57
+
58
+ interface.launch(debug=True)