← back to Handbag Authentication

handbag_data/github_datasets/Pursearch/src/app.py

425 lines

import tensorflow as tf
import torch
import torch.nn as nn
from torchvision import models
import torchvision.transforms as transforms
from sklearn.metrics.pairwise import cosine_similarity
from PIL import Image
from werkzeug.utils import secure_filename
import pathlib

import io
import numpy as np
import pandas as pd
import os
import sys

from flask import Flask, jsonify, request, flash, redirect, url_for, render_template, send_from_directory

app = Flask(__name__)

# Define dirs and requirements
total_brands = ['Louis Vuitton', 'Chanel', 'Coach', 'Michael Kors', 'Gucci', 'Longchamp', 'Other', 'None']
brands = [total_brands[i] for i in [1,2,0,6,1,5,4,4,1,3,3,0,0,7,0,3,4,0,4]]
names = ['2.55','F57842', 'alma','handbag','boy','classic','dionysus','ggmarmont','grandshoppingtote','jetsettote','mercertotebag',
'neverfull','newwavechaintote','None','petiteboitechapeau','selmasatchel','sohodisco','speedy','sylvie']
websites = ["https://www.chanel.com/en_WW/fashion.html#search/2.55",
"https://www.coach.com/",
"https://us.louisvuitton.com/eng-us/search/alma",
"None",
"https://www.chanel.com/en_WW/fashion.html#search/BOY/",
"https://us.longchamp.com/products/le-pliage/tote-bag-s/l2605089545",
"https://www.gucci.com/us/en/st/capsule/dionysus-collection",
"https://www.gucci.com/us/en/st/capsule/gg-marmont-collection",
"https://www.chanel.com/us/fashion/p/A91046Y09958N5024/large-shopping-bag-calfskin-silver-tone-metal/",
"https://www.michaelkors.com/search/_/Rtt-jet+set+",
"https://www.michaelkors.com/search/_/Rtt-mercer+",
"https://us.louisvuitton.com/eng-us/search/neverfull",
"https://us.louisvuitton.com/eng-us/search/new%20wave%20chain",
"None",
"https://us.louisvuitton.com/eng-us/products/petite-boite-chapeau-monogram-nvprod670013v",
"https://www.michaelkors.com/search/_/Rtt-selma+satchel",
"https://www.gucci.com/us/en/st/newsearchpage?searchString=soho%20disco&search-cat=header-search",
"https://us.louisvuitton.com/eng-us/search/speedy",
"https://www.gucci.com/us/en/st/newsearchpage?searchString=sylvie&search-cat=header-search"
]

UPLOAD_FOLDER = '../data/example'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0
ALLOWED_EXTENSIONS = { 'png', 'jpg', 'jpeg'}
tf.gfile = tf.io.gfile

# Model loading and image preprocessing (crop)
def load_detection_model(model_name):
    """
    Load detection model from tensorflow
    
    Parameters
    ----------
    model_name : str
        model name that will be loaded, can be found in 
        /Users/jianinglu1/Documents/GitHub/models/research/object_detection/g3doc/detection_model_zoo.md

    Returns
    -------
    model
        detection model
    """
    base_url = 'http://download.tensorflow.org/models/object_detection/'
    model_file = model_name + '.tar.gz'
    model_dir = tf.keras.utils.get_file(
        fname=model_name, 
        origin=base_url + model_file,
        untar=True)

    model_dir = pathlib.Path(model_dir)/"saved_model"

    model = tf.saved_model.load(str(model_dir))
    model = model.signatures['serving_default']

    return model

def load_model():
    """
    Load the identification model for handbags
    
    Returns
    -------
    model_ft
        whole model for predict brand and type
    model_test
        model for get the handbag vector (remove the last layer of whole model)
    """
    model_ft = models.resnet50(pretrained=True)
    num_ftrs = model_ft.fc.in_features
    # Here the size of each output sample is set to 2.
    # Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
    model_ft.fc = nn.Linear(num_ftrs, 19)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model_ft = model_ft.to(device)
    model_name = torch.load("../models/model_name_resnet50gpu_correct_merged", map_location=torch.device('cpu'))
    model_ft.load_state_dict(model_name)
    model_ft.eval()
    # get the model for extracting the handbag vectors
    model_test = nn.Sequential(*list(model_ft.children())[:-1])
    model_test.eval()
    return model_ft, model_test

def run_inference_for_single_image(model, image):
    """
    Get object detection results
    
    Parameters
    ----------
    model : object
        detection model
    image : image object
        input image
    
    Returns
    -------
    dictionary
        object detection results
    """
    image = np.asarray(image)
    # The input needs to be a tensor, convert it using `tf.convert_to_tensor`.
    input_tensor = tf.convert_to_tensor(image)
    # The model expects a batch of images, so add an axis with `tf.newaxis`.
    input_tensor = input_tensor[tf.newaxis,...]

    # Run inference
    output_dict = model(input_tensor)

    # All outputs are batches tensors.
    # Convert to numpy arrays, and take index [0] to remove the batch dimension.
    # We're only interested in the first num_detections.
    num_detections = int(output_dict.pop('num_detections'))
    output_dict = {key:value[0, :num_detections].numpy() 
                 for key,value in output_dict.items()}
    output_dict['num_detections'] = num_detections

    # detection_classes should be ints.
    output_dict['detection_classes'] = output_dict['detection_classes'].astype(np.int64)

    return output_dict

def crop_image(image_bytes, model):
    """
    Crop image based on object detection results and only keep the handbag part of image
    
    Parameters
    ----------
    image_bytes : input image, read from input image file
        input image
    model : object
        detection model
    
    Returns
    -------
    img
        crop img if handbag detection is successfult, None if no handbag has been detected.
    """
    image = Image.open(io.BytesIO(image_bytes))
    width, hight = image.size
    image_np = np.array(image)
    output_dict = run_inference_for_single_image(model, image_np)
    if 31 in output_dict["detection_classes"]:
        max_score =  max([output_dict["detection_scores"][idx] for idx, i in enumerate(output_dict["detection_classes"]) if i == 31])
        index = list(output_dict["detection_scores"]).index(max_score) 
        [y,x,h,w] = output_dict["detection_boxes"][index]
        cropped = image_np[int(y*hight):int(h*hight), int(x*width):int(w*width)]
        img = Image.fromarray(cropped,"RGB")
        return img
    else:
        return False

def transform_image(image):
    """
    Transform image to get the input for identification model
    
    Parameters
    ----------
    image : image object
        input image
    
    Returns
    -------
    transformed image
    """
    my_transforms = transforms.Compose([transforms.Resize((256,256)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    return my_transforms(image).unsqueeze(0)

def get_prediction(image, model, model_vector):
    """
    Get identification results
    
    Parameters
    ----------
    image : image object
        transformed image
    model : object
        identification model
    model_vector : object
        model used to get handbag vector
    
    Returns
    -------
    prediction results
    """
    tensor = transform_image(image=image)
    outputs = model(tensor)
    vectors = model_vector(tensor).detach().numpy()
    vectors = np.reshape(vectors, [1,2048])
    _, y_hat = outputs.max(1)
    predicted_idx = y_hat.item()

    return brands[predicted_idx], names[predicted_idx], websites[predicted_idx], vectors


# Load identification model 
model, model_vector = load_model()
# Load detection model
# faster rcnn trained using coco 
model_name = 'faster_rcnn_inception_v2_coco_2018_01_28'
detection_model = load_detection_model(model_name)

# Handbag recommendatation 
# read handbag information data
shop_data_shein = pd.read_csv("../data/shop/file/shop_shein_index_final.csv")
shop_data_saks = pd.read_csv("../data/shop/file/shop_bag_index_final.csv")
vectors_shop = torch.load("../data/shop/file/shop_bag_vect_reorder.pt")
vectors_shein = torch.load("../data/shop/file/shop_shein_bag_vect_reorder.pt")
vectors_shop = np.reshape(vectors_shop,[2344,2048])
vectors_shein = np.reshape(vectors_shein,[1270,2048])

def recommend_bag_sim(input_vectors, vectors_bags):
    """
    Calculate similarity of handbags and rank handbags
    
    Parameters
    ----------
    input_vectors : numpy array
        input handbag vector
    vectors_bags : numpy array
        retailer handbags vector
    
    Returns
    -------
    sim_data
        data frame with data ranked by similarity
    """
    sim = cosine_similarity(input_vectors, vectors_bags)
    sim = np.reshape(sim, np.shape(sim)[1])
    sim_dict = [["sim_" + str(i),j] for (i,j) in zip(range(len(sim)), sim)]
    
    sim_data = pd.DataFrame(sim_dict)
    sim_data[2] = [ i.split("_")[0] for i in list(sim_data[0])]
    sim_data["order"] = sim_data[1].rank(ascending=False)
    sim_data.columns = ["initial_index", "sim", "index", "order"]
    return sim_data

def recommend_bag(input_vector, shop_data_saks, shop_data_shein, price_range, vetors_shein, vectors_saks, company):
    """
    Recommend Top3 handbags based criteria
    
    Parameters
    ----------
    input_vector : numpy array
        input handbag vector
    shop_data_saks : data frame
        saks data set
    shop_data_shein : data frame
        shein data set
    price_range : str
        price range required by user
    vetors_shein : numpy array
        handbag vectors from shein
    vectors_shop : numpy array
        handbag vectors from saks
    company : str
        retailer required by user
    
    Returns
    -------
    Top3 recommended handbags
    """
    price_dict = {"less 50": ["less 50"],
                 "less 250" : ["less 50", "less 250"],
                 "less 500" : ["less 50", "less 250", "less 500"],
                 "less 750" : ["less 50", "less 250", "less 500", "less 750"],
                 "less 1000" : ["less 50", "less 250", "less 500", "less 750", "less 1000"],
                 "all" : ["less 50", "less 250", "less 500", "less 750", "less 1000", "higher 1000"]}
    if price_range == "less 50" or company == "shein":
        sim_data = recommend_bag_sim(input_vector, vectors_shein)
        shop_data = pd.merge(shop_data_shein, sim_data, on = shop_data_shein.index)
        shop_data.drop("key_0", axis=1, inplace=True)
        shop_data = shop_data[shop_data["price_level"].isin(price_dict[price_range])].sort_values("sim", ascending=False)
        #Top5 = shop_data[0:5]["image"].values
    
    elif company == "saks":
        sim_data = recommend_bag_sim(input_vector, vectors_shop)
        shop_data = pd.merge(shop_data_saks, sim_data, on = shop_data_saks.index)
        shop_data.drop("key_0", axis=1, inplace=True)
        shop_data = shop_data[shop_data["price_level"].isin(price_dict[price_range])].sort_values("sim", ascending=False)
        #Top5 = shop_data[0:5]["image"].values
    
    else:
        sim_data_saks = recommend_bag_sim(input_vector, vectors_shop)
        shop_data_saks = pd.merge(shop_data_saks, sim_data_saks, on = shop_data_saks.index)
        shop_data_saks.drop("key_0", axis=1, inplace=True)
        
        sim_data_shein = recommend_bag_sim(input_vector, vectors_shein)
        shop_data_shein = pd.merge(shop_data_shein, sim_data_shein, on = shop_data_shein.index)
        shop_data_shein.drop("key_0", axis=1, inplace=True)
        
        shop_data = pd.concat([shop_data_saks, shop_data_shein], axis=0)
        shop_data = shop_data[shop_data["price_level"].isin(price_dict[price_range])].sort_values("sim", ascending=False)
        #Top5 = shop_data[0:5]["image"].values
    Top_1 = get_infor(shop_data, 0)
    Top_2 = get_infor(shop_data, 1)
    Top_3 = get_infor(shop_data, 2)
    return Top_1, Top_2, Top_3

def get_infor(shop_data, index):
    """
    Get the information of data
    
    Parameters
    ----------
    shop_data : data frame
        handbag data set
    index : int
        index of rows
    
    Returns
    -------
    handbag information
    """
    shop_data = shop_data[index:index+1]
    return [shop_data["page"].values[0], shop_data["product"].values[0], shop_data["brand"].values[0], shop_data["price"].values[0], shop_data["image"].values[0], round(shop_data["sim"].values[0],2)]

# Build web app 
@app.route("/")
def hello():
    return render_template('home.html')

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route('/uploads/<filename>', methods=['GET', 'POST'])
def uploaded_file(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)

@app.route('/predict', methods=['GET', 'POST'])
def predict():
    if request.method == 'POST':
        if "file" not in request.files:
            return redirect(request.url)
        test = request.form["test"]
        if test == "test":
            file = open(os.path.join(app.config['UPLOAD_FOLDER'], "test.jpg"), 'rb')
            filename = "test.jpg"
        else:
            file = request.files['file']
            if not file:
                return render_template("no_result.html", message="No Image Uploaded!")

            if file and allowed_file(file.filename):
                filename = secure_filename(file.filename)
                savename = os.path.join(app.config["UPLOAD_FOLDER"], filename)
                if os.path.exists(savename):
                    os.remove(savename)
                file.save(savename)
                file.stream.seek(0)
                
        img_bytes = file.read()
        option = request.form.get('options')
        if option == "yes":
            infile = crop_image(img_bytes, detection_model)
            if infile == False:
                return render_template("no_result.html", message="Handbag Detection Failed, Retry with No Crop!")
            savename_crop = os.path.join(app.config["UPLOAD_FOLDER"], filename.split(".")[0] + "_crop.jpg")
            if os.path.exists(savename_crop):
                os.remove(savename_crop)
            infile.save(os.path.join(app.config["UPLOAD_FOLDER"], filename.split(".")[0] + "_crop.jpg"))
            crop_filename = filename.split(".")[0] + "_crop.jpg"
        else:
            infile = Image.open(io.BytesIO(img_bytes))
            crop_filename = filename
        
        brand, name, web, input_vector = get_prediction(image=infile, model=model, model_vector=model_vector)
        if name != "None":
            price_range = request.form.get("price")
            company = request.form.get("company")
            Top1, Top2, Top3 = recommend_bag(input_vector, shop_data_saks, shop_data_shein, price_range, vectors_shein, vectors_shop, company)
            return render_template('result.html', class_brand=brand, class_name=name, class_web=web,
            Top1_product=Top1[1], Top2_product=Top2[1], Top3_product=Top3[1],
            Top1_sim=Top1[-1], Top2_sim=Top2[-1], Top3_sim=Top3[-1],
            Top1_page=Top1[0], Top2_page=Top2[0], Top3_page=Top3[0],
            Top1_brand=Top1[2], Top2_brand=Top2[2], Top3_brand=Top3[2],
            Top1_price=Top1[3], Top2_price=Top2[3], Top3_price=Top3[3],
            Top1_img=Top1[4],Top2_img=Top2[4], Top3_img=Top3[4],
            filename=filename,
            crop_filename=crop_filename)
        else:
            return render_template("no_result.html", message="No Handbag in Uploaded Image!")
    return render_template('index.html')

@app.after_request
def add_header(response):
    # response.cache_control.no_store = True
    response.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, post-check=0, pre-check=0, max-age=0'
    response.headers['Pragma'] = 'no-cache'
    response.headers['Expires'] = '-1'
    return response

if __name__ == '__main__':
    print("Loading PyTorch model and Flask starting server ...")
    print("Please wait until server has fully started")
    app.run(host="0.0.0.0", port=80, debug=False)
    ### Command to run this: FLASK_ENV=development FLASK_APP=app.py FLASK_DEBUG=0 flask run