Spaces:
Sleeping
Sleeping
| 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() | |