Spaces:
Sleeping
Sleeping
File size: 2,805 Bytes
2710430 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
from flask import Flask, request, render_template, jsonify
from huggingface_hub import hf_hub_download
import cv2
import numpy as np
import tensorflow as tf
import pandas as pd
import base64
import json
# Initialize Flask app
app = Flask(__name__)
# Load model and data at startup
# model = tf.keras.models.load_model("sneaker_category_predictor_v2.h5")
model_path = hf_hub_download(repo_id="thejagstudio/SneakerAI", filename="sneaker_category_predictor_v2.h5")
model = tf.keras.models.load_model(model_path)
# Define expected columns for one-hot encoding
with open("metadata.json", "r") as f:
METADATA_COLUMNS = json.load(f)
def encode_metadata(data):
# Create DataFrame with single row
df = pd.DataFrame({k: [v.lower()] for k, v in data.items()})
# Initialize empty DataFrame with all possible columns
encoded = pd.DataFrame()
# Encode each feature maintaining consistent columns
for feature, possible_values in METADATA_COLUMNS.items():
feature_encoded = pd.get_dummies(df[feature], prefix=feature)
# Add missing columns with 0s
for value in possible_values:
col_name = f"{feature}_{value}"
if col_name not in feature_encoded.columns:
feature_encoded[col_name] = 0
encoded = pd.concat([encoded, feature_encoded], axis=1)
# Ensure consistent column order
all_columns = [
f"{feat}_{val}" for feat, vals in METADATA_COLUMNS.items() for val in vals
]
encoded = encoded.reindex(columns=all_columns, fill_value=0)
return encoded.values.astype(np.float32)
@app.route("/")
def index():
global METADATA_COLUMNS
return render_template("index.html", metadata=METADATA_COLUMNS)
@app.route("/predict", methods=["POST"])
def predict():
try:
data = request.json
# Process image
img_data = base64.b64decode(data["image"])
img_array = np.frombuffer(img_data, np.uint8)
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
img = cv2.resize(img, (224, 224))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img / 255.0
img = np.expand_dims(img, axis=0)
# Encode metadata with consistent columns
metadata = encode_metadata(
{
"brand": data["brand"],
"color": data["color"],
"gender": data["gender"],
"midsole": data["midsole"],
"upperMaterial": data["upperMaterial"],
}
)
# Make prediction
predictions = model.predict([img, metadata])
categories = [
"Lifestyle",
"Running",
"Other",
"Cleat",
"Sandal",
"Basketball",
"Boot",
"Skateboarding",
]
confidenceList = predictions[0].tolist()
return jsonify({"categories": categories, "confidence": confidenceList})
except Exception as e:
return jsonify({"error": str(e)}), 400
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)
|