-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy path0_tag_folder.py
216 lines (181 loc) · 9.03 KB
/
0_tag_folder.py
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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
import os
import requests
import yaml
from pathlib import Path
import os
import numpy as np
from PIL import Image
import shutil
import subprocess
import uuid
import argparse
import yaml
from advcaption.taggers import ImageTagger
import json
# from advcaption.nlp import load_model as loadmodel, loadinflectmodel, process_input_tags
# Specify the directory containing the image files
IMAGE_DIR = 'H:/cannabis/leafly'
API_ENDPOINT = 'http://127.0.0.1:9999/tag-image/'
def copy_template_to_image_directory(image_directory, template_file='./templates/template.yaml'):
"""
Copies the template yaml file './templates/template.yaml' into the specified image directory.
Parameters:
- image_directory: The path to the directory where the template YAML file will be copied.
"""
if not os.path.isdir(image_directory):
print(f"The directory {image_directory} does not exist.")
return {}
if not os.path.isfile(template_file):
print("Template file './templates/template.yaml' does not exist.")
return {}
destination_path = os.path.join(image_directory, os.path.basename(template_file))
if os.path.isfile(destination_path):
print(f"Template file already exists at {destination_path}.")
try:
with open(destination_path, 'r') as file:
template_data = yaml.safe_load(file)
print("Template YAML loaded into dictionary.")
return template_data
except yaml.YAMLError as e:
print(f"Error loading YAML from template file: {e}")
return {}
shutil.copy(template_file, destination_path)
print(f"Template file copied to {destination_path}.")
try:
with open(destination_path, 'r') as file:
template_data = yaml.safe_load(file)
print("Template YAML loaded into dictionary.")
return template_data
except yaml.YAMLError as e:
print(f"Error loading YAML from template file: {e}")
return {}
def load_image(image_path, error_directory):
try:
image = Image.open(image_path) # Load the image using PIL
img = np.array(image, np.uint8) # Attempt the conversion to numpy array
except OSError as e:
print(f"Error: {e}. Moving to error directory.")
if not os.path.exists(error_directory): # Create error directory if it doesn't exist
os.makedirs(error_directory)
shutil.move(image_path, os.path.join(error_directory, os.path.basename(image_path))) # Move the erroneous image
return None
return img
def validate_images_imagemagick(image_directory, error_directory):
if not os.path.exists(error_directory): # Create the error directory if it doesn't exist
os.makedirs(error_directory)
for filename in os.listdir(image_directory):
# Update the condition to include .webp and .webm
if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff', '.webp', '.webm')):
image_path = os.path.join(image_directory, filename)
try:
subprocess.run(['identify', image_path], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # Validate image
except subprocess.CalledProcessError as e:
print(f"Error occurred with file {filename}: {e}")
shutil.move(image_path, os.path.join(error_directory, filename)) # Move the erroneous image
def rename_files_in_directory(directory, rename=True):
"""
Recursively renames all image files in the specified directory and its subdirectories.
Each subdirectory will have its files renamed to a unique pattern followed by an incrementing number.
Parameters:
- directory: The root directory containing the files to be renamed.
"""
if not os.path.isdir(directory):
print(f"The directory {directory} does not exist.")
return
new_files = []
for dirpath, dirnames, filenames in os.walk(directory):
# Generate a unique pattern for each subdirectory
pattern = generate_unique_id(length=5)
counter = 1
# Sort files for consistent renaming across runs
filenames.sort()
for filename in filenames:
if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff', '.webp', '.webm')):
if rename:
extension = os.path.splitext(filename)[1]
new_filename = f"{pattern}_{counter}{extension}"
old_file = os.path.join(dirpath, filename)
new_file = os.path.join(dirpath, new_filename)
os.rename(old_file, new_file)
new_files.append(new_file)
counter += 1
else:
old_file = os.path.join(dirpath, filename)
new_files.append(old_file)
return new_files
def generate_unique_id(length=8):
"""
Generates a unique identifier using a substring of a UUID4.
Parameters:
- length: The desired length of the identifier. Default is 8 characters.
Returns:
A string representing the unique identifier.
"""
# Generate a random UUID
unique_id_full = uuid.uuid4()
# Convert the UUID to a string, remove hyphens, and take a substring of the specified length
unique_id_short = str(unique_id_full).replace('-', '')[:length]
return unique_id_short
def send_image_and_save_response(image_path: Path, url: str):
image_path = Path(image_path)
yaml_path = image_path.with_suffix('.yaml')
if yaml_path.exists():
print(f"{yaml_path.name} already exists. Skipping reprocessing.")
return # Exit the function if the YAML file exists
with open(image_path, 'rb') as image_file:
# Prepare the 'file' part of the multipart/form-data request
files = {'file': (image_path.name, image_file, 'image/jpeg')}
# Make the POST request to the FastAPI endpoint
response = requests.post(url, files=files)
if response.status_code == 200:
# Convert the response to YAML format
response_yaml = yaml.dump(response.json(), allow_unicode=True, sort_keys=False)
# Define the path for the YAML file (same name as the image, but with .yaml extension)
yaml_path = image_path.with_suffix('.yaml')
# Write the YAML data to the file
with open(yaml_path, 'w') as yaml_file:
yaml_file.write(response_yaml)
else:
print(f"Error with {image_path.name}: {response.text}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Process images in a directory.")
parser.add_argument("--image_directory", help="The directory containing images to process.")
parser.add_argument("--error_directory", help="The directory where error images will be moved.")
parser.add_argument("--beast", action="store_true", help="Use z3d Tagger for furry concepts.")
parser.add_argument("--rename", action="store_true", help="Enable renaming of error images.")
parser.add_argument("--overwrite", action="store_true", help="Overwrite the tags from a previous run.")
args = parser.parse_args()
# Validate Part 1
for dirpath, dirnames, filenames in os.walk(args.image_directory):
for filename in filenames:
# Check if the file extension indicates an image file
if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff', '.webp', '.webm')):
image_path = os.path.join(dirpath, filename)
_ = load_image(image_path, args.error_directory)
renamed_filelist = rename_files_in_directory(args.image_directory, rename=args.rename)
template_file = os.path.join(args.image_directory, './templates.yaml')
if os.path.exists(template_file):
try:
with open(template_file, 'r', encoding='utf-8') as file:
template_data = yaml.safe_load(file)
print("Template YAML loaded into dictionary.")
config = template_data
except Exception as e:
print(f"Failed to load template: {e}")
config = {}
else:
config = copy_template_to_image_directory(args.image_directory)
for filename in renamed_filelist:
print(filename)
# Check if a corresponding YAML file already exists
yaml_path = os.path.splitext(filename)[0] + '.yaml'
if os.path.exists(yaml_path) and args.overwrite == False:
print(f"YAML file already exists for {filename}. Skipping.")
continue
# Proceed with processing the image if no YAML file exists
if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff', '.webp', '.webm')):
try:
send_image_and_save_response(filename, API_ENDPOINT)
print(f"Saved combined results to {yaml_path}.")
except Exception as e: # Catching a general exception to handle any kind of failure in process_image
print(f"Failed to process image {filename}. Error: {e}")