dd

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+import requests
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
+
+
+def greet(url):
+
+    processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-large-cityscapes-semantic")
+    model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-large-cityscapes-semantic")
+
+    image = Image.open(requests.get(url, stream=True).raw)
+    inputs = processor(images=image, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+    # model predicts class_queries_logits of shape `(batch_size, num_queries)`
+    # and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
+    class_queries_logits = outputs.class_queries_logits
+    masks_queries_logits = outputs.masks_queries_logits
+
+    # you can pass them to processor for postprocessing
+    predicted_semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
+
+    color_map = {
+        0: (0, 0, 0),  # 클래스 0: 검은색
+        1: (255, 255, 255),  # 클래스 1: 흰색
+        2: (255, 0, 0),
+        3: (0, 255, 0),
+        4: (0, 0, 255),
+        5: (255, 255, 0),
+        6: (255, 0, 255),
+        7: (0, 255, 255),
+        # 다른 클래스에 대한 색상 지정
+    }
+    #semantic_image = Image.new('RGB', predicted_semantic_map.shape[1:][::-1])[0]
+    semantic_image = Image.new('RGB', (predicted_semantic_map.shape[1], predicted_semantic_map.shape[0]))
+    pixels = semantic_image.load()
+    for y in range(predicted_semantic_map.shape[0]):
+        for x in range(predicted_semantic_map.shape[1]):
+            class_id = predicted_semantic_map[y, x].item()
+            color = color_map.get(class_id, (0, 0, 0))
+            pixels[x, y] = color
+
+    return pixels
+
+
+url = "https://www.cityscapes-dataset.com/wordpress/wp-content/uploads/2015/07/muenster00.png"
+greet(url)
+
+iface = gr.Interface(
+    fn=greet,
+    inputs=gr.Image(value=url),
+    outputs="image"
+)
+
+iface.launch(share=True)