Spaces:
Sleeping
Sleeping
File size: 3,933 Bytes
56ac15e | 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 | import gradio as gr
import os
import shutil
import cv2
import numpy as np
from naming import im2c
from collections import Counter
COLOR_NAME = ['black', 'brown', 'blue', 'gray', 'green', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
def get_top_names(img):
# resize images to smaller size
anchor = 512
width = img.shape[1]
height = img.shape[0]
if width > 1024 or height > 1024:
if width >= height:
dim = (np.floor(width/height*anchor).astype(int), anchor)
else:
dim = (anchor, np.floor(height/width*anchor).astype(int))
img = cv2.resize(img, dim, interpolation=cv2.INTER_LINEAR)
w2c = np.load('w2c11_j.npy').astype(np.float16)
_, _, name_idx_img, _ = im2c(img, w2c)
filtered_counts = Counter(name_idx_img[name_idx_img <= 10])
sorted_counts = sorted(filtered_counts.items(), key=lambda x: x[1], reverse=True)
top_3_values = [num for num, count in sorted_counts[:3]]
top_3_colors = [COLOR_NAME[i] for i in top_3_values]
# print("Top 3 colors:", top_3_colors)
return top_3_values, top_3_colors
def classify_and_log(images):
output_folder = "classified_results"
os.makedirs(output_folder, exist_ok=True)
category_folders = {i: os.path.join(output_folder, COLOR_NAME[i]) for i in range(11)}
for folder in category_folders.values():
os.makedirs(folder, exist_ok=True)
log_file = os.path.join(output_folder, "log.txt")
results = {i: [] for i in range(11)}
with open(log_file, "w") as log:
for id_img, img in enumerate(images):
filename = os.path.basename(img.name)
img_array = cv2.imread(img).astype(np.float32)
img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
cat_id, category = get_top_names(img_array)
# print(category_folders[cat_id[0]], filename)
target_path = os.path.join(category_folders[cat_id[0]], filename)
cv2.imwrite(target_path, cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB))
print(f"Image:{filename} -> Top 3 colors:{category}\n")
log.write(f"{filename} -> {category}\n")
results[cat_id[0]].append(target_path)
# output_images = [results[i] for i in range(11)]
# print(output_images)
return log_file
def swap_to_gallery(images):
return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
def upload_example_to_gallery(images, prompt, style, negative_prompt):
return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
def remove_back_to_files():
return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
with gr.Blocks() as demo:
gr.Markdown("## Image color categorization")
with gr.Row():
with gr.Column():
image_input = gr.File(
label="Drag (Select) more than one photos",
file_types=["image"],
file_count="multiple"
)
uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)
with gr.Column(visible=False) as clear_button:
remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=image_input, size="sm")
image_input.upload(fn=swap_to_gallery, inputs=image_input, outputs=[uploaded_files, clear_button, image_input])
remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, image_input])
classify_btn = gr.Button("submit")
# with gr.Row():
# image_output = {str(i): gr.Gallery(label=f"{i}") for i in range(11)}
log_output = gr.File(label="download logs")
classify_btn.click(
classify_and_log,
inputs=[image_input],
outputs=[log_output]
)
demo.launch()
|