Update app.py

app.py

@@ -1,3 +1,23 @@
+import gradio as gr
+import torch
+from qai_hub_models.models.detr_resnet50 import Model
+from PIL import Image, ImageDraw
+import numpy as np
+
+# 注册 AVIF 支持（根据所安装的插件选择一种）
+try:
+    from pillow_avif import register_avif_opener
+    register_avif_opener()
+except ImportError:
+    try:
+        import pillow_heif
+        pillow_heif.register_heif_opener()
+    except ImportError:
+        print("AVIF support not available. Please install 'pillow-avif-plugin' or 'pillow-heif'.")
+
+# 加载模型
+torch_model = Model.from_pretrained()
+
 def detect_objects(image):
     if image is None:
         raise ValueError("No image uploaded!")  # 检查图像是否为 None
@@ -38,10 +58,40 @@ def detect_objects(image):
                 "box": box,
             })
             
-            # 检查并绘制边界框和标签到原始图像上
+            # 绘制边界框和标签到原始图像上
             draw = ImageDraw.Draw(original_image)
             if box[1] < box[3]:  # 确保 y_min < y_max
                 draw.rectangle(box[:4], outline="red", width=3)  
                 draw.text((box[0], box[1]), f"{detections[-1]['label']} ({detections[-1]['confidence']})", fill="red")  
 
     return original_image, detections
+
+# 创建 Gradio 接口，自动处理 AVIF 图像并转换为 PNG 格式以供显示和处理。
+with gr.Blocks() as iface:
+    gr.Markdown("# Object Detection with DETR-ResNet50")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="numpy", label="Upload Image (supports PNG, JPEG, AVIF...)")  
+            submit_button = gr.Button("Submit")
+            clear_button = gr.Button("Clear")
+        with gr.Column(scale=1):
+            output_image = gr.Image(label="Detected Image")
+            output_json = gr.JSON(label="Detections")
+
+    def on_submit(image):
+        try:
+            detected_image, detections = detect_objects(image)
+            return detected_image, detections
+        except Exception as e:
+            return None, {"error": str(e)}
+
+    def on_clear():
+        return None, None, None  # 清空输入和输出
+
+    submit_button.click(on_submit, inputs=image_input, outputs=[output_image, output_json])
+    clear_button.click(on_clear, inputs=None, outputs=[image_input, output_image, output_json])  # 修复清除功能
+
+# 启动应用程序
+if __name__ == "__main__":
+    iface.launch()