1

app.py

@@ -1,35 +1,41 @@
 import gradio as gr
 from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
 from PIL import Image
+import numpy as np
 import torch
 
 # 모델과 feature extractor 로드
-model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b3-finetuned-cityscapes-1024-1024")
-feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b3-finetuned-cityscapes-1024-1024")
+model_name = "nvidia/segformer-b0-finetuned-ade-512-512"
+model = SegformerForSemanticSegmentation.from_pretrained(model_name)
+feature_extractor = SegformerFeatureExtractor.from_pretrained(model_name)
 
-# 이미지를 처리하는 함수
-def predict(image):
-    # 이미지를 모델에 맞게 변환
-    processed_image = image.resize((1024, 1024))
-    inputs = feature_extractor(images=processed_image, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
+def segment_image(image):
+    # 이미지 처리
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-    # 결과를 이미지로 변환 (예: 가장 높은 확률을 가진 클래스 선택)
-    result = torch.argmax(logits)
-    result = result.squeeze().cpu().numpy()
+    # 마스크 생성
+    upsampled_logits = torch.nn.functional.interpolate(
+        outputs.logits, size=image.size[::-1], mode="bilinear", align_corners=False
+    )
+    upsampled_predictions = upsampled_logits.argmax(dim=1)
+    mask = upsampled_predictions.squeeze().numpy()
 
-    # 여기에서는 단순화를 위해 결과를 그대로 반환합니다.
-    # 실제로는 결과를 적절한 형식으로 변환해야 할 수 있습니다.
-    return result
+    # 결과 반환
+    return Image.fromarray(np.uint8(mask * 255))
 
-# Gradio 인터페이스 생성
+# 예시 이미지 경로
+example_images = ["image1.jpg", "image2.jpg", "image3.jpg"]
+
+# Gradio 인터페이스 설정
 demo = gr.Interface(
-    fn=predict,
-    inputs=gr.inputs.Image(type='pil'),
-    outputs=gr.outputs.Image(type='pil'),
-    examples=["image1.jpg", "image2.jpg", "image3.jpg"]  # 세 개의 예제 이미지 경로
+    fn=segment_image,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs="image",
+    title="머신러닝 7주차 과제_3",
+    examples=example_images
 )
 
 # 인터페이스 실행
-demo.launch()
+demo.launch()