aliabd HF staff commited on
Commit
ce6eb81
1 Parent(s): d074d90

Delete app.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. app.py +0 -42
app.py DELETED
@@ -1,42 +0,0 @@
1
- import gradio as gr
2
- import torch
3
- import random
4
- import numpy as np
5
- from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
6
-
7
- device = torch.device("cpu")
8
- model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-ade").to(device)
9
- model.eval()
10
- preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-ade")
11
-
12
- def visualize_instance_seg_mask(mask):
13
- image = np.zeros((mask.shape[0], mask.shape[1], 3))
14
- labels = np.unique(mask)
15
- label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}
16
- for i in range(image.shape[0]):
17
- for j in range(image.shape[1]):
18
- image[i, j, :] = label2color[mask[i, j]]
19
- image = image / 255
20
- return image
21
-
22
- def query_image(img):
23
- target_size = (img.shape[0], img.shape[1])
24
- inputs = preprocessor(images=img, return_tensors="pt")
25
- with torch.no_grad():
26
- outputs = model(**inputs)
27
- outputs.class_queries_logits = outputs.class_queries_logits.cpu()
28
- outputs.masks_queries_logits = outputs.masks_queries_logits.cpu()
29
- results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0].cpu().detach()
30
- results = torch.argmax(results, dim=0).numpy()
31
- results = visualize_instance_seg_mask(results)
32
- return results
33
-
34
- demo = gr.Interface(
35
- query_image,
36
- inputs=[gr.Image()],
37
- outputs="image",
38
- title="MaskFormer Demo",
39
- examples=[["example_2.png"]]
40
- )
41
-
42
- demo.launch()