1

app.py

@@ -1,18 +1,34 @@
+import gradio as gr
 from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
 from PIL import Image
+import torch
 
+# 모델과 feature extractor 로드
+model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 
-feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b3-finetuned-cityscapes-1024-1024")
-model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b3-finetuned-cityscapes-1024-1024")
+# 이미지를 처리하는 함수
+def predict(image):
+    # 이미지를 모델에 맞게 변환
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits
 
-# url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-# image = Image.open(requests.get(url, stream=True).raw)
+    # 결과를 이미지로 변환 (예: 가장 높은 확률을 가진 클래스 선택)
+    result = torch.argmax(logits, dim=1)
+    result = result.squeeze().cpu().numpy()
 
-image1, image2, image3 = 'image1.jpg', 'image2.jpg', 'image3.jpg'
-image1 = Image.open(image1)
-image2 = Image.open(image2)
-image3 = Image.open(image3)
+    # 여기에서는 단순화를 위해 결과를 그대로 반환합니다.
+    # 실제로는 결과를 적절한 형식으로 변환해야 할 수 있습니다.
+    return result
 
-inputs = feature_extractor(images=[image1, image2, image3], return_tensors="pt")
-outputs = model(**inputs)
-logits = outputs.logits  # shape (batch_size, num_labels, height/4, width/4)
+# Gradio 인터페이스 생성
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(shape=(400, 600)),
+    outputs=gr.outputs.Image(),
+    examples=["image1.jpg", "image2.jpg", "image3.jpg"]  # 세 개의 예제 이미지 경로
+)
+
+# 인터페이스 실행
+demo.launch()