elliesleightholm
commited on
Commit
·
d09d79e
1
Parent(s):
d3576a5
adding basic classification structure
Browse files- .gitignore +4 -0
- app.py +142 -0
- images/bike-helmet.png +0 -0
- images/coffee.png +0 -0
- images/cutting-board.png +0 -0
- images/flip-flops.jpg +0 -0
- images/grater.png +0 -0
- images/iron.png +0 -0
- images/laptop.png +0 -0
- images/notebook.png +0 -0
- images/sleeping-bag.png +0 -0
- items.py +46 -0
- requirements.txt +2 -0
.gitignore
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
.env
|
2 |
+
__pycache__
|
3 |
+
venv
|
4 |
+
.DS_Store
|
app.py
ADDED
@@ -0,0 +1,142 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import open_clip
|
3 |
+
import torch
|
4 |
+
import requests
|
5 |
+
import numpy as np
|
6 |
+
from PIL import Image
|
7 |
+
from io import BytesIO
|
8 |
+
from items import ecommerce_items
|
9 |
+
import os
|
10 |
+
from dotenv import load_dotenv
|
11 |
+
|
12 |
+
# Load environment variables from the .env file
|
13 |
+
load_dotenv()
|
14 |
+
|
15 |
+
# Sidebar content
|
16 |
+
sidebar_markdown = """
|
17 |
+
|
18 |
+
Note, this demo can classify 200 items. If you didn't find what you're looking for, reach out to us on our [Community](https://join.slack.com/t/marqo-community/shared_invite/zt-2iab0260n-QJrZLUSOJYUifVxf964Gdw) and request an item to be added.
|
19 |
+
|
20 |
+
## Documentation
|
21 |
+
📚 [Blog Post]()
|
22 |
+
|
23 |
+
📝 [Use Case Blog Post]()
|
24 |
+
|
25 |
+
## Code
|
26 |
+
💻 [GitHub Repo]()
|
27 |
+
|
28 |
+
🤝 [Google Colab]()
|
29 |
+
|
30 |
+
🤗 [Hugging Face Collection]()
|
31 |
+
|
32 |
+
## Citation
|
33 |
+
If you use Marqo-Ecommerce-L or Marqo-Ecommerce-B, please cite us:
|
34 |
+
```
|
35 |
+
|
36 |
+
```
|
37 |
+
"""
|
38 |
+
|
39 |
+
from huggingface_hub import login
|
40 |
+
|
41 |
+
# Get your Hugging Face API key (ensure it is set in your environment variables)
|
42 |
+
api_key = os.getenv("HF_API_TOKEN")
|
43 |
+
|
44 |
+
if api_key is None:
|
45 |
+
raise ValueError("Hugging Face API key not found. Please set the 'HF_API_TOKEN' environment variable.")
|
46 |
+
|
47 |
+
# Login using the token
|
48 |
+
login(token=api_key)
|
49 |
+
|
50 |
+
# Initialize the model and tokenizer
|
51 |
+
def load_model(progress=gr.Progress()):
|
52 |
+
progress(0, "Initializing model...")
|
53 |
+
model_name = 'hf-hub:Marqo/marqo-ecommerce-embeddings-B'
|
54 |
+
model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(model_name)
|
55 |
+
|
56 |
+
progress(0.5, "Loading tokenizer...")
|
57 |
+
tokenizer = open_clip.get_tokenizer(model_name)
|
58 |
+
|
59 |
+
text = tokenizer(ecommerce_items)
|
60 |
+
|
61 |
+
progress(0.75, "Encoding text features...")
|
62 |
+
with torch.no_grad(), torch.amp.autocast('cuda'):
|
63 |
+
text_features = model.encode_text(text)
|
64 |
+
text_features /= text_features.norm(dim=-1, keepdim=True)
|
65 |
+
|
66 |
+
progress(1.0, "Model loaded successfully!")
|
67 |
+
|
68 |
+
return model, preprocess_val, text_features
|
69 |
+
|
70 |
+
# Load model and prepare interface
|
71 |
+
model, preprocess_val, text_features = load_model()
|
72 |
+
|
73 |
+
# Prediction function
|
74 |
+
def predict(image, url):
|
75 |
+
if url:
|
76 |
+
response = requests.get(url)
|
77 |
+
image = Image.open(BytesIO(response.content))
|
78 |
+
|
79 |
+
processed_image = preprocess_val(image).unsqueeze(0)
|
80 |
+
|
81 |
+
with torch.no_grad(), torch.amp.autocast('cuda'):
|
82 |
+
image_features = model.encode_image(processed_image)
|
83 |
+
image_features /= image_features.norm(dim=-1, keepdim=True)
|
84 |
+
|
85 |
+
text_probs = (100 * image_features @ text_features.T).softmax(dim=-1)
|
86 |
+
|
87 |
+
sorted_confidences = sorted(
|
88 |
+
{ecommerce_items[i]: float(text_probs[0, i]) for i in range(len(ecommerce_items))}.items(),
|
89 |
+
key=lambda x: x[1],
|
90 |
+
reverse=True
|
91 |
+
)
|
92 |
+
top_10_confidences = dict(sorted_confidences[:10])
|
93 |
+
|
94 |
+
return image, top_10_confidences
|
95 |
+
|
96 |
+
# Clear function
|
97 |
+
def clear_fields():
|
98 |
+
return None, ""
|
99 |
+
|
100 |
+
# Gradio interface
|
101 |
+
title = "Ecommerce Item Classifier with Marqo-Ecommerce Embedding Models"
|
102 |
+
description = "Upload an image or provide a URL of a fashion item to classify it using Marqo-Ecommerce Models!"
|
103 |
+
|
104 |
+
examples = [
|
105 |
+
["images/laptop.png", "Laptop"],
|
106 |
+
["images/grater.png", "Grater"],
|
107 |
+
["images/flip-flops.jpg", "Flip Flops"],
|
108 |
+
["images/bike-helmet.png", "Bike Helmet"],
|
109 |
+
["images/sleeping-bag.png", "Sleeping Bag"],
|
110 |
+
["images/cutting-board.png", "Cutting Board"],
|
111 |
+
["images/iron.png", "Iron"],
|
112 |
+
["images/coffee.png", "Coffee"],
|
113 |
+
]
|
114 |
+
|
115 |
+
with gr.Blocks(css="""
|
116 |
+
.remove-btn {
|
117 |
+
font-size: 24px !important; /* Increase the font size of the cross button */
|
118 |
+
line-height: 24px !important;
|
119 |
+
width: 30px !important; /* Increase the width */
|
120 |
+
height: 30px !important; /* Increase the height */
|
121 |
+
}
|
122 |
+
""") as demo:
|
123 |
+
with gr.Row():
|
124 |
+
with gr.Column(scale=1):
|
125 |
+
gr.Markdown(f"# {title}")
|
126 |
+
gr.Markdown(description)
|
127 |
+
gr.Markdown(sidebar_markdown)
|
128 |
+
gr.Markdown(" ", elem_id="vertical-line") # Add an empty Markdown with a custom ID
|
129 |
+
with gr.Column(scale=2):
|
130 |
+
input_image = gr.Image(type="pil", label="Upload Fashion Item Image", height=312)
|
131 |
+
input_url = gr.Textbox(label="Or provide an image URL")
|
132 |
+
with gr.Row():
|
133 |
+
predict_button = gr.Button("Classify")
|
134 |
+
clear_button = gr.Button("Clear")
|
135 |
+
gr.Markdown("Or click on one of the images below to classify it:")
|
136 |
+
gr.Examples(examples=examples, inputs=input_image)
|
137 |
+
output_label = gr.Label(num_top_classes=6)
|
138 |
+
predict_button.click(predict, inputs=[input_image, input_url], outputs=[input_image, output_label])
|
139 |
+
clear_button.click(clear_fields, outputs=[input_image, input_url])
|
140 |
+
|
141 |
+
# Launch the interface
|
142 |
+
demo.launch()
|
images/bike-helmet.png
ADDED
images/coffee.png
ADDED
images/cutting-board.png
ADDED
images/flip-flops.jpg
ADDED
images/grater.png
ADDED
images/iron.png
ADDED
images/laptop.png
ADDED
images/notebook.png
ADDED
images/sleeping-bag.png
ADDED
items.py
ADDED
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
ecommerce_items = [
|
2 |
+
"laptop", "smartphone", "tablet", "headphones", "smartwatch", "camera",
|
3 |
+
"printer", "gaming console", "TV", "keyboard", "mouse", "monitor",
|
4 |
+
"external hard drive", "USB drive", "memory card", "router", "modem",
|
5 |
+
"vacuum cleaner", "washing machine", "refrigerator", "microwave",
|
6 |
+
"air conditioner", "toaster", "blender", "coffee maker", "kitchen appliance",
|
7 |
+
"furniture", "mattress", "bedding", "shoes", "clothing", "handbag",
|
8 |
+
"watch", "sunglasses", "cosmetics", "skincare", "perfume", "books",
|
9 |
+
"toys", "video games", "jewelry", "fitness equipment", "bike", "sports gear",
|
10 |
+
"pet supplies", "office supplies", "stationery", "gardening tools",
|
11 |
+
"power tools", "home decor", "art supplies", "musical instruments",
|
12 |
+
"baby products", "diapers", "groceries", "snacks", "health supplements",
|
13 |
+
"vitamins", "medications", "cleaning supplies", "water bottle", "backpack",
|
14 |
+
"notebook", "planner", "pen", "pencil", "water filter", "electric kettle",
|
15 |
+
"hair dryer", "hair straightener", "shaving kit", "toothbrush", "toothpaste",
|
16 |
+
"soap", "shampoo", "conditioner", "razor", "deodorant", "body lotion",
|
17 |
+
"hand sanitizer", "first aid kit", "flashlight", "batteries", "extension cord",
|
18 |
+
"light bulb", "fan", "heater", "humidifier", "dehumidifier", "iron",
|
19 |
+
"clothes steamer", "vacuum bags", "toolbox", "drill", "screwdriver",
|
20 |
+
"hammer", "nails", "tape measure", "wrench", "pliers", "ladder", "paint",
|
21 |
+
"paintbrush", "roller", "drop cloth", "caulk", "sealant", "screws", "bolts",
|
22 |
+
"nuts", "faucet", "showerhead", "toilet paper", "paper towels", "trash bags",
|
23 |
+
"cleaning cloth", "disinfectant wipes", "mop", "broom", "dustpan",
|
24 |
+
"scrubbing brush", "laundry detergent", "fabric softener", "dryer sheets",
|
25 |
+
"dish soap", "sponge", "cutlery", "plates", "bowls", "cups", "glasses",
|
26 |
+
"wine glasses", "baking sheets", "pots", "pans", "cutting board",
|
27 |
+
"kitchen knives", "grater", "peeler", "mixing bowls", "measuring cups",
|
28 |
+
"measuring spoons", "oven mitts", "apron", "spatula", "whisk", "can opener",
|
29 |
+
"wine opener", "dish rack", "salt", "pepper", "herbs", "spices", "oil",
|
30 |
+
"vinegar", "sugar", "flour", "rice", "pasta", "cereal", "bread",
|
31 |
+
"cookies", "crackers", "chips", "coffee", "tea", "juice", "soda",
|
32 |
+
"water", "wine", "beer", "whiskey", "vodka", "rum", "champagne",
|
33 |
+
"cocktail shaker", "ice tray", "pizza cutter", "cheese slicer",
|
34 |
+
"tupperware", "food storage bags", "aluminum foil", "plastic wrap",
|
35 |
+
"parchment paper", "disposable plates", "disposable cups", "plastic utensils",
|
36 |
+
"grill", "BBQ tools", "charcoal", "firewood", "cooler", "picnic basket",
|
37 |
+
"lawn chair", "outdoor furniture", "hammock", "tent", "sleeping bag",
|
38 |
+
"camping stove", "lantern", "fishing rod", "bicycle helmet", "running shoes",
|
39 |
+
"yoga mat", "dumbbells", "jump rope", "resistance bands", "treadmill",
|
40 |
+
"elliptical machine", "stationary bike", "protein powder", "energy bars",
|
41 |
+
"water bottle", "goggles", "swimsuit", "sunscreen", "insect repellent",
|
42 |
+
"beach towel", "sandals", "flip flops", "suitcase", "travel pillow",
|
43 |
+
"passport holder", "luggage tag", "camera tripod", "memory card reader",
|
44 |
+
"portable charger", "phone case", "screen protector", "sunhat", "hat", "trousers",
|
45 |
+
"t-shirt", "shirt", "socks", "shoes",
|
46 |
+
]
|
requirements.txt
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
open_clip_torch
|
2 |
+
transformers
|