Update app.py

app.py

@@ -1,72 +1,45 @@
-from PIL import Image
-from yolo import YOLO
 import gradio as gr
-import os
+import torch
+from torchvision import models, transforms
+from PIL import Image
 
-# Initialize YOLO model
-yolo = YOLO()
+# 定义设备
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-def detect_objects(image, crop=False, count=True):
-    r_image = yolo.detect_image(image, crop=crop, count=count)
-    return r_image
+# 加载预训练的模型
+model = models.resnet18(pretrained=True)
+model = model.to(device)
+model.eval()
 
-def save_image(image, filename):
-    if not os.path.exists("img_out"):
-        os.makedirs("img_out")
-    image.save(os.path.join("img_out", filename), quality=95, subsampling=0)
-    return os.path.join("img_out", filename)
+# 图像预处理
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
 
-# Gradio interface for single image prediction
-def predict(image):
-    result_image = detect_objects(image)
-    output_path = save_image(result_image, "output.png")
-    return result_image
+# 加载类名称
+with open("model_data/rtts_classes.txt") as f:
+    class_names = [line.strip() for line in f.readlines()]
 
-# Function to list images in the 'img' folder
-def get_image_list():
-    image_folder = "img"
-    if not os.path.exists(image_folder):
-        os.makedirs(image_folder)
-    img_names = [f for f in os.listdir(image_folder) if f.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff'))]
-    img_paths = [os.path.join(image_folder, img_name) for img_name in img_names]
-    return img_paths
-
-# Gradio interface components
-image_output = gr.Image(type="pil", label="Output Image")
-image_select = gr.inputs.Radio(get_image_list, label="Select Image from 'img' Folder")
+# 定义预测函数
+def predict(image):
+    image = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        outputs = model(image)
+    _, predicted = outputs.max(1)
+    return class_names[predicted]
 
-# Gradio app for single image prediction
+# 使用Gradio创建界面
 iface = gr.Interface(
-    fn=predict, 
-    inputs=image_select, 
-    outputs=image_output,
-    title="YOLO Object Detection",
-    description="Select an image from the 'img' folder to detect objects using YOLO model."
+    fn=predict,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=gr.outputs.Textbox(),
+    title="图像分类器",
+    description="上传一张图像，并让模型预测它的类别。",
 )
 
-# Custom function to display images in a gallery format
-def gallery_update():
-    image_list = get_image_list()
-    return gr.Gallery.update(value=image_list)
-
-# Gradio gallery component for displaying images in the 'img' folder
-gallery = gr.Gallery(get_image_list(), label="Image Gallery")
-
-# Update gallery when Gradio interface loads
-gallery.update(gallery_update)
-
-# Combine both interfaces
-app = gr.Blocks()
-
-with app:
-    with gr.Row():
-        with gr.Column():
-            gallery.render()
-        with gr.Column():
-            image = gr.Image(type="pil", label="Selected Image")
-            predict_button = gr.Button("Predict")
-            output_image = gr.Image(type="pil", label="Output Image")
-            predict_button.click(predict, inputs=image, outputs=output_image)
-
-# Launch the app
-app.launch()
+# 启动应用
+if __name__ == "__main__":
+    iface.launch()