dd

app.py

@@ -3,6 +3,7 @@ import requests
 import torch
 from PIL import Image
 from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
+import numpy as np
 
 
 def greet(url):
@@ -21,29 +22,17 @@ def greet(url):
     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
+    predicted_semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[url.size])[0]
+
+    sepia_filter = np.array([
+        [0.393, 0.769, 0.189],
+        [0.349, 0.686, 0.168],
+        [0.272, 0.534, 0.131]
+    ])
+    sepia_img = predicted_semantic_map.dot(sepia_filter.T)
+    sepia_img /= sepia_img.max()
+
+    return sepia_img
 
 
 url = "http://images.cocodataset.org/val2017/000000039769.jpg"
@@ -55,4 +44,4 @@ iface = gr.Interface(
     outputs="image"
 )
 
-iface.launch()
+iface.launch(debug = True)