使用Flask部署图像分类模型

sudo apt-get install python3-flask

!pip install torch torchvision

# 导入所需的库
import json
import io
import glob
from PIL import Image
from torchvision import models
import torchvision.transforms as transforms

# 将参数“pretraining”传递为“True”，使用预训练的权重:
model = models.densenet121(pretrained=True)
# 切换到模型到“eval”模式:
model.eval()

# 定义预处理的函数
def transform_image(image_bytes):
    my_transforms = transforms.Compose([transforms.Resize(255),
                                        transforms.CenterCrop(224),
                                        transforms.ToTensor(),
                                        transforms.Normalize(
                                            [0.485, 0.456, 0.406],
                                            [0.229, 0.224, 0.225])])
    image = Image.open(io.BytesIO(image_bytes))
    return my_transforms(image).unsqueeze(0)

# 加载由pytorch提供的映射
imagenet_class_mapping = json.load(open('imagenet_class_index.json'))

# 定义函数来获得图片的预测
# 它接受参数:图片路径并提供预测作为输出
def get_category(image_path):
  #以二进制形式读取图像
    with open(image_path, 'rb') as file:
        image_bytes = file.read()
    # 变换图像
    transformed_image = transform_image(image_bytes=image_bytes)
    # 使用模型来预测类
    outputs = model.forward(transformed_image)
    _, category = outputs.max(1)
    # 返回
    predicted_idx = str(category.item())
    return imagenet_class_mapping[predicted_idx]

get_category(image_path='static/sample_1.jpeg')
## ['n02089973', 'English_foxhound']

get_category(image_path='static/sample_2.jpeg')
## ['n11939491', 'daisy']

# 导入所需的库
import requests
from bs4 import BeautifulSoup
import os
import time

def get_path(url):
    return "static/URL_" + str(url.replace("/","_"))
  
headers = {
    'User-Agent': "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/77.0.3865.90 Safari/537.36"
    } 

# 定义爬取图像并将其存储在目录中的函数
def get_images(url):
   # get the directory path
    path = get_path(url)
    try:
        os.mkdir(path)
    except:
        pass
    # 从URL请求源代码
    response = requests.request("GET", url, headers=headers)
    # 通过Beautiful Soup解析数据
    data = BeautifulSoup(response.text, 'html.parser')
    # 在源代码中找到图像标记
    images = data.find_all('img', src=True)
    # 从所有的图像标签中提取src
    image_src = [x['src'] for x in images]
    # 只选择jpeg
    image_src = [x for x in image_src if x.endswith('.jpeg') ]
    image_count = 1
    # 在指定目录存储图像
    for image in image_src:
        print(image)
        image_file_name = path+'/'+str(image_count)+'.jpeg' 
        print(image_file_name)
        # 以写入二进制形式打开文件并添加图像内容来存储它
        with open(image_file_name, 'wb') as f:
            res = requests.get(image)
            f.write(res.content)
        image_count = image_count+1

get_images('https://medium.com/@allanishac/9-wild-animals-that-would-make-a-much-better-president-than-donald-trump-b41f960bb171')

def get_picture_html(path, tag):
    image_html = """<p> {tag_name} </p> <picture> <img src= "../{path_name}"  height="300" width="400"> </picture>"""
    return image_html.format(tag_name=tag, path_name=path)

# 定义在html文件中添加列表元素的函数
def get_count_html(category, count):
    count_html = """<li> {category_name} : {count_} </li>"""
    return count_html.format(category_name = category, count_ = count)

# 计数
def get_value_count(image_class_dict):
    count_dic = {}
    for category in image_class_dict.values():
        if category in count_dic.keys():
            count_dic[category] = count_dic[category]+1
        else:
            count_dic[category] = 1
    return count_dic

# 函数从image_class字典生成html文件
# 键将是图像的路径，而值将是与之关联的类。
def generate_html(image_class_dict):
    picture_html = ""
    count_html = ""
    
    # 循环这些键并将图像添加到html文件中
    for image in image_class_dict.keys():
        picture_html += get_picture_html(path=image, tag= image_class_dict[image])
        
    value_counts = get_value_count(image_class_dict)
    
    # 循环value_counts并向html文件中添加类的计数
    for value in value_counts.keys():
        count_html += get_count_html(value, value_counts[value])

# 导入库
from flask import Flask, render_template, request, redirect, url_for
from get_images import get_images, get_path, get_directory
from get_prediction import get_prediction
from generate_html import generate_html
from torchvision import models
import json

app = Flask(__name__)

# 映射
imagenet_class_mapping = json.load(open('imagenet_class_index.json'))

# 使用预训练模型
model = models.densenet121(pretrained=True)
model.eval()

# 定义从url获取图像并预测类的函数
def get_image_class(path):
    # 从URL获取图像并将其存储在给定的路径中
    get_images(path)
    # 根据所提供的目录预测图像的图像类别
    path = get_path(path)
    images_with_tags = get_prediction(model, imagenet_class_mapping, path)
    # 生成html文件以在我们预测类之后呈现
    generate_html(images_with_tags)

# 根页面为"home.html"    
@app.route('/')
def home():
    return render_template('home.html')

@app.route('/', methods=['POST', 'GET'])
def get_data():
    if request.method == 'POST':
        user = request.form['search']
        # 如果搜索按钮被点击，调用函数get_image_class
        get_image_class(user)
        #返回image_class.html
        return redirect(url_for('success', name=get_directory(user)))


@app.route('/success/<name>')
def success(name):
    return render_template('image_class.html')


if __name__ == '__main__' :
    app.run(debug=True)

# 获取目录中出现的所有图像的类
def get_category(model, imagenet_class_mapping, image_path):
    with open(image_path, 'rb') as file:
        image_bytes = file.read()
    transformed_image = transform_image(image_bytes=image_bytes)
    outputs = model.forward(transformed_image)
    _, category = outputs.max(1)
    
    predicted_idx = str(category.item())
    return imagenet_class_mapping[predicted_idx]

# 它将创建一个图像路径和预测类的字典
# 我们将使用该字典生成html文件。
def get_prediction(model, imagenet_class_mapping, path_to_directory):
    files = glob.glob(path_to_directory+'/*')
    image_with_tags = {}
    for image_file in files:
        image_with_tags[image_file] = get_category(model, imagenet_class_mapping, image_path=image_file)[1]
    return image_with_tags