Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
# importing the required libaries
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__)
# mapping
imagenet_class_mapping = json.load(open('imagenet_class_index.json'))
# use the pre-trained model
model = models.densenet121(pretrained=True)
# define the function to get the images from the url and predicted the class
def get_image_class(path):
# get images from the URL and store it in a given path
# predict the image class of the images with provided directory
path = get_path(path)
images_with_tags = get_prediction(model, imagenet_class_mapping, path)
# generate html file to render once we predict the classes
# by deafult render the "home.html"
def home():
return render_template('home.html')
@app.route('/', methods=['POST', 'GET'])
def get_data():
if request.method == 'POST':
user = request.form['search']
# if search button hit, call the function get_image_class
#render the image_class.html
return redirect(url_for('success', name=get_directory(user)))
def success(name):
return render_template('image_class.html')
if __name__ == '__main__' :
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment