nielsr HF staff commited on
Commit
a5fde5f
1 Parent(s): ea301e4

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -0
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
2
+ import gradio as gr
3
+ import torch
4
+ import matplotlib.pyplot as plt
5
+ import cv2
6
+
7
+ processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
8
+ model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
9
+
10
+ def process_image(image, prompts):
11
+ inputs = processor(text=prompts, images=[image] * len(prompts), padding="max_length", return_tensors="pt")
12
+
13
+ # predict
14
+ with torch.no_grad():
15
+ outputs = model(**inputs)
16
+ preds = outputs.logits.unsqueeze(1)
17
+
18
+ filename = f"mask.png"
19
+ plt.imsave(filename,torch.sigmoid(preds[1][0]))
20
+
21
+ img2 = cv2.imread(filename)
22
+ gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
23
+
24
+ (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
25
+
26
+ # fix color format
27
+ cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
28
+
29
+ return Image.fromarray(bw_image)
30
+
31
+ title = "Interactive demo: zero-shot image segmentation with CLIPSeg"
32
+ description = "Demo for using CLIPSeg, a CLIP-based model for zero- and one-shot image segmentation. 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."
33
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.10003'>CLIPSeg: Image Segmentation Using Text and Image Prompts</a> | <a href='https://huggingface.co/docs/transformers/main/en/model_doc/clipseg'>HuggingFace docs</a></p>"
34
+
35
+ examples = [["a glass", "something to fill", "wood", "a jar"]]
36
+
37
+ interface = gr.Interface(fn=process_image,
38
+ inputs=[gr.Image(type="pil"), gr.Textbox(label="What do you want to identify (separated by comma)?")],
39
+ outputs=gr.Image(type="pil"),
40
+ title=title,
41
+ description=description,
42
+ article=article,
43
+ examples=examples)
44
+
45
+ interface.launch(debug=True)