Spaces:
Running
Running
File size: 7,575 Bytes
e267b1d f01686c 58a3046 25cc30d e267b1d 5aaefae f01686c b1e655a 6fec7e7 7a8c805 1fb0946 5aaefae 0764de0 5aaefae eef83e3 76e01ce 5aaefae 76e01ce 5aaefae f9f191a 1fb0946 76e01ce 5aaefae 76e01ce 56dafcf 76e01ce 56dafcf 76e01ce 56dafcf 76e01ce 56dafcf 76e01ce 56dafcf 5201feb f9f191a 1fb0946 f9f191a e267b1d 5aaefae e267b1d e45b47d a69b8f5 7a8c805 5aaefae b6666d0 5aaefae 1fb0946 9718868 1fb0946 5aaefae 76e01ce f01686c 5aaefae 76e01ce 5aaefae 76e01ce 5aaefae d54d724 76e01ce 5aaefae 76e01ce 5aaefae 76e01ce 5aaefae 76e01ce 5aaefae 76e01ce b6666d0 1fb0946 e9ff897 97ece73 e9ff897 7a8c805 97ece73 1fb0946 b6666d0 76e01ce f01686c 1fb0946 d54d724 76e01ce 5aaefae bb3c537 f01686c 76e01ce f01686c 76e01ce f01686c 76e01ce 1fb0946 b6666d0 1fb0946 76e01ce 5aaefae |
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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
import streamlit as st
from ultralytics import YOLO
from PIL import Image
import time
from transformers import AutoModelForImageClassification, AutoImageProcessor
import requests
from io import BytesIO
import numpy as np
import cv2
import concurrent.futures
from classification import pipes
from rembg import remove
from PIL import Image
# input_path = 'input.png'
# output_path = 'output.png'
# input = Image.open(input_path)
# output = remove(input)
# output.save(output_path)
time_taken={}
# Categories dictionary
categories_dict = {
"UpperBody": ["shirt, blouse","top, t-shirt, sweatshirt", "blouse", "sweater", "cardigan", "jacket", "vest"],
"Lowerbody": ["pants", "shorts", "skirt"],
"Wholebody": ["coat", "dress", "jumpsuit", "cape"],
"Head": ["glasses", "hat", "headband", "head covering", "hair accessory"],
"Neck": ["tie", "neckline", "collar"],
"Arms and Hands": ["glove", "watch", "sleeve"],
"Waist": ["belt"],
"Legs and Feet": ["leg warmer", "tights", "stockings", "sock", "shoe"],
"Others": ["bag", "wallet", "scarf", "umbrella"],
"Garment parts": ["hood", "lapel", "epaulette", "pocket"],
"Closures": ["buckle", "zipper"],
"Decorations": ["applique", "bead", "bow", "flower", "fringe", "ribbon", "rivet", "ruffle", "sequin", "tassel"]
}
def find_category(subcategory):
for category, subcategories in categories_dict.items():
if subcategory in subcategories:
return category
return "Subcategory not found."
# Load models and processor only once using Streamlit session state
if 'models_loaded' not in st.session_state:
sm=time.time()
# localization model
st.session_state.segment_model = YOLO("best.pt")
# image preprocessor
st.session_state.image_processor = AutoImageProcessor.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-topwear")
# top wear
st.session_state.top_wear_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-topwear")
# bottom wear
st.session_state.bottomwear_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-bottomwear")
# full wear
st.session_state.fullwear = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-fullwear")
# for full wear and top wear
st.session_state.pattern_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-pattern-rgb")
st.session_state.print_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-print")
st.session_state.sleeve_length_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-sleeve-length")
st.session_state.neck_style_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-neck-style")
st.session_state.models_loaded = True
time_taken["model loading"]=time.time()-sm
# Streamlit app UI
st.title("Clothing Classification Pipeline")
url = st.sidebar.text_input("Paste image URL here...")
if url:
try:
response = requests.get(url)
if response.status_code == 200:
image = Image.open(BytesIO(response.content))
st.sidebar.image(image.resize((200, 200)), caption="Uploaded Image", use_column_width=False)
# Ensure the image is in RGBA mode (or RGB)
if image.mode != 'RGB':
image = image.convert('RGB')
rmimg = remove(image)
# Convert image to numpy array for YOLO model
image_np = np.array(image)
outputs={}
# Perform inference
ss=time.time()
# results = st.session_state.segment_model(image_np)
results = st.session_state.segment_model(image_np, conf=0.55)
time_taken["yolo model result time"]=time.time()-ss
# Create a copy of the original image to draw bounding boxes and labels
output_image = image_np.copy()
cropped_images_list = [] # List to hold cropped images and their titles
# Visualize the segmentation results
for result in results:
boxes = result.boxes # Bounding boxes
classes = result.names # Class names of the detected objects
for i, box in enumerate(boxes):
box_coords = box.xyxy[0].cpu().numpy().astype(int)
x1, y1, x2, y2 = box_coords
# Draw the bounding box on the original image
cv2.rectangle(output_image, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=1)
# Get the class label and confidence score for the object
class_label = classes[box.cls[0].int().item()]
confidence = box.conf[0].item()
# Prepare the label text with class and confidence
label_text = f'{class_label}: {confidence:.2f}'
# Put text label on the original image
cv2.putText(output_image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# Crop the image based on the bounding box
cropped_image = image_np[y1:y2, x1:x2].copy()
cropped_image_resized = cv2.resize(cropped_image, (200, 200)) # Resize cropped image
category_name = find_category(class_label)
# st.write(f"Detected category: {category_name}")
cs=time.time()
if category_name=="Neck":
st.write("neck")
outputs[category_name]=pipes(cropped_image, category_name)
else:
outputs[category_name]=pipes(rmimg, category_name)
st.write("other")
time_taken[f"{category_name} prediction time"]=time.time()-cs
# st.write(pipes(cropped_image, category_name))
# Add cropped image and its title to the list
cropped_images_list.append((cropped_image_resized, f'Class: {category_name}, Confidence: {confidence:.2f}'+">>subcat>>"+label_text))
time_taken["whole process time"]=time.time()-ss
# Display the original image with bounding boxes and labels
st.sidebar.image(output_image, caption="Segmented Image", channels="RGB", use_column_width=True)
# Display cropped images row-wise
num_columns = 3 # Number of columns per row
num_rows = (len(cropped_images_list) + num_columns - 1) // num_columns # Calculate the number of rows
for i in range(num_rows):
cols = st.columns(num_columns)
for j in range(num_columns):
idx = i * num_columns + j
if idx < len(cropped_images_list):
cropped_image, title = cropped_images_list[idx]
with cols[j]:
st.image(cropped_image, caption=title, use_column_width=True)
st.header("Output")
st.json(outputs)
st.header("taken time")
st.json(time_taken)
else:
st.write("URL Invalid...!")
except Exception as e:
st.write(f"An error occurred: {e}")
|