Update app.py

app.py

@@ -1,17 +1,91 @@
+# import gradio as gr
+# from transformers import pipeline
+
+# pipe = pipeline("image-classification", model="chaojie777/google-vit-base-patch16-224-in21k")
+# images = ["images/daisy.jpg", "images/dandelion.jpg", "images/rosa.jpg", "images/sunflower.jpg", "images/tulip.jpg"]
+
+# iface = gr.Interface.from_pipeline(
+#         pipe, 
+#         examples= [ [example] for example in images],
+#         description="Final project that labels flowers images into: Daisy, Dandelion, Rose, Sunflower, Tulip",
+#         title="Flower Classifier - Vit"
+# )
+
+# iface.launch()
+import torch
+import torchvision.transforms as transforms
+from torchvision import models
+from torch import nn
+import torch.nn.functional as F
 import gradio as gr
-from transformers import pipeline
-
-pipe = pipeline("image-classification", model="chaojie777/google-vit-base-patch16-224-in21k")
+from PIL import Image
+import json
+
+# Get cpu or gpu device for training.
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using {device} device")
+
+# 设置模型路径
+model_path = './best.pth'  # 替换为您训练的模型的路径
+num_classes = 5
+label_name_list = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
+
+# 创建并加载模型
+resnet18 = models.resnet18(pretrained=True)
+num_ftrs = resnet18.fc.in_features
+resnet18.fc = nn.Linear(num_ftrs, num_classes)
+resnet18 = resnet18.to(device)
+
+# 加载训练好的模型参数
+resnet18.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+model = resnet18
+model.eval()
+
+# 图片转换
+train_transform = transforms.Compose([
+    # transforms.RandomRotation(5),
+    # transforms.RandomHorizontalFlip(),
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+def predict(inp):
+    # 定义预处理变换
+    transform = train_transform
+
+    # 加载图片并进行预处理
+    image = transform(inp).unsqueeze(0).to(device)
+
+    # 使用模型进行预测
+    with torch.no_grad():
+        output = model(image)
+
+    # 数据后处理
+    # 计算预测概率
+    pred_score = nn.functional.softmax(output[0], dim=0)
+    pred_score = pred_score.cpu().numpy()
+
+    # 获取预测结果
+    pred_index = torch.argmax(output, dim=1).item()
+    pred_label = label_name_list[pred_index]
+
+    # 转为json字符串格式
+    result_dict = {'pred_score': str(max(pred_score)), 'pred_index': str(pred_index), 'pred_label': pred_label}
+    result_json = json.dumps(result_dict)
+
+    return result_json
 images = ["images/daisy.jpg", "images/dandelion.jpg", "images/rosa.jpg", "images/sunflower.jpg", "images/tulip.jpg"]
+# 设置Gradio接口
+demo = gr.Interface(fn=predict,
+                    inputs=gr.Image(type="pil"),
+                    outputs="text",
+                    examples= [ [example] for example in images],
+                    )
 
-iface = gr.Interface.from_pipeline(
-        pipe, 
-        examples= [ [example] for example in images],
-        description="Final project that labels flowers images into: Daisy, Dandelion, Rose, Sunflower, Tulip",
-        title="Flower Classifier - Vit"
-)
+# 启动Gradio接口
+demo.launch()
 
-iface.launch()
 # import torch
 # import torchvision.transforms as transforms
 # from torchvision import models