Update app.py

return to gradio

app.py

@@ -1,28 +1,21 @@
-from flask import Flask, request, jsonify
-from flask_cors import CORS
+import gradio as gr
 import cv2
 import numpy as np
-import matplotlib
-matplotlib.use('Agg')  # 非GUI后端
+
+
 import matplotlib.pyplot as plt
 import time
 import tempfile
-from sklearn.cluster import MiniBatchKMeans
-from PIL import Image
 from transformers import ViTImageProcessor, ViTForImageClassification
 import torch
 import concurrent.futures
-import io
-import os
+import io # Import io for handling byte streams
 
-# 加载预训练的 ViT 模型和图像处理器
+# Load the pre-trained Vision Transformer model and image processor
 model_name = "google/vit-base-patch16-224"
 image_processor = ViTImageProcessor.from_pretrained(model_name)
 model = ViTForImageClassification.from_pretrained(model_name)
 
-app = Flask(__name__)
-CORS(app)  # 启用 CORS
-
 def create_color_palette(colors, size=100):
     palette = np.zeros((4*size, 4*size, 3), dtype=np.uint8)
     for i in range(4):
@@ -47,7 +40,7 @@ def identify_image(image):
     
     return ", ".join(results)
 
-def analyze_image(image_data):
+def analyze_image(image_data): # Modified to accept image data
     """Analyzes the uploaded image."""
     # Load the image data into a PIL Image object
     image = Image.open(io.BytesIO(image_data))
@@ -118,25 +111,24 @@ def analyze_image(image_data):
 
     return top_5_classes, histogram_path, edges_path, palette_path, timing_output
 
-@app.route('/spawn', methods=['POST'])
-def spawn():
-    # 获取上传的图像数据
-    image_data = request.files['image'].read()
-
-    # 调用分析图像的函数
-    top_5_classes, histogram_path, edges_path, palette_path, timing_output = analyze_image(image_data)
-
-    # 构建响应数据
-    response_data = {
-        'top_5_classes': top_5_classes,
-        'histogram_path': histogram_path,
-        'edges_path': edges_path,
-        'palette_path': palette_path,
-        'timing_output': timing_output
-    }
-
-    # 返回 JSON 响应
-    return jsonify(response_data)
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8080)
+# Define Gradio interface
+interface = gr.Interface(
+    fn=analyze_image,
+    inputs=gr.Image(type="numpy", label="Upload Image"),
+    outputs=[
+        gr.Textbox(label="Image Classification (Top 5)"),
+        gr.Image(label="Histogram"),
+        gr.Image(label="Edge Detection"),
+        gr.Image(label="Main Colors Palette"),
+        gr.Textbox(label="Processing Times")
+    ],
+    title="Advanced Image Analysis",
+    description="Upload an image to analyze its classification, histogram, edges, and main colors.",
+    allow_flagging="never"
+)
+
+# Set request queue parameters
+interface.queue(max_size=5)
+
+# Launch the interface
+interface.launch()