Update app.py

app.py

@@ -1,31 +1,32 @@
 import gradio as gr
 from transformers import DetrForObjectDetection, DetrFeatureExtractor
+from PIL import Image, ImageDraw
 import torch
 
-# 모델 ID 지정
 model_id = "facebook/detr-resnet-50"
-
-# 모델과 피처 추출기 로드
 model = DetrForObjectDetection.from_pretrained(model_id)
 feature_extractor = DetrFeatureExtractor.from_pretrained(model_id)
 
 def object_detection(image):
-    # 이미지를 모델 입력으로 변환
     inputs = feature_extractor(images=image, return_tensors="pt")
     outputs = model(**inputs)
 
-    # 결과 처리
-    # 여기서는 간단히 모델의 출력 logits를 리스트로 반환합니다.
-    # 실제 애플리케이션에서는 출력을 사용자가 이해할 수 있는 형태로 가공해야 합니다.
-    return outputs.logits.tolist()
+    # 후처리 과정
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = feature_extractor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+
+    draw = ImageDraw.Draw(image)
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [round(i, 2) for i in box.tolist()]
+        draw.rectangle(box, outline="red", width=3)
+        draw.text((box[0], box[1]), f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}", fill="red")
+
+    return image
 
-# Gradio 인터페이스 정의 및 실행
 iface = gr.Interface(
     fn=object_detection,
-    inputs=gr.Image(),
-    outputs="json",
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
     title="Object Detection with DETR",
-    description="Upload an image and the model will detect objects in the image."
-)
-
-iface.launch()
+    description="Upload an image, and the model will detect objects in the image. Detected objects will be highlighted with bounding boxes."
+).launch()