creek
The AI Image Editor of 2030git clone https://github.com/maydayv7/creek
models/texture/texture.py (14743B)
1# -*- coding: utf-8 -*-
2"""
3Texture and Material
4Original file: https://colab.research.google.com/drive/1u-5cHHckvZPMjIFJ5YJcCAGzDWVPRl91
5"""
6
7## .pth -> .pt
8
9import numpy as np
10import pandas as pd
11import torch
12import torch.nn as nn
13import numpy as np
14import os
15
16CENTROIDS_PATH = "texture_centroids_augmented.pth_path"
17OUTPUT_PATH = "texture_centroids_augmented.pt_path"
18
19
20class DataWrapper(nn.Module):
21 def __init__(self, data_dict):
22 super().__init__()
23 self.data = torch.jit.Attribute(data_dict, dict[str, torch.Tensor])
24
25 def forward(self):
26 return self.data
27
28
29if os.path.exists(CENTROIDS_PATH):
30 print("Loading .pth file...")
31 raw_data = torch.load(CENTROIDS_PATH, map_location="cpu", weights_only=False)
32
33 print("Processing data...")
34 clean_dict = {}
35 for k, v in raw_data.items():
36 if isinstance(v, np.ndarray):
37 clean_dict[k] = torch.from_numpy(v).float()
38 else:
39 clean_dict[k] = v.float()
40
41 wrapper = DataWrapper(clean_dict)
42 scripted_wrapper = torch.jit.script(wrapper)
43 scripted_wrapper.save(OUTPUT_PATH)
44 print(f"Success! Saved as {OUTPUT_PATH}")
45else:
46 print("File not found.")
47
48## Generate centroids
49
50import torch
51import torchvision.transforms as T
52from PIL import Image
53import numpy as np
54import os
55import glob
56from typing import Dict, List
57
58DATASET_ROOT = "Texture_Dataset_path"
59OUTPUT_PATH = "texture_centroids_s14.pt"
60DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
61
62
63print(f"Loading DINOv2-S14 on {DEVICE}...")
64model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
65model.to(DEVICE)
66model.eval()
67
68transform = T.Compose(
69 [
70 T.Resize(256, interpolation=T.InterpolationMode.BICUBIC),
71 T.CenterCrop(224),
72 T.ToTensor(),
73 T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
74 ]
75)
76
77
78def get_folder_centroid(folder_path):
79 image_files = glob.glob(os.path.join(folder_path, "*"))
80 valid_exts = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
81 image_files = [f for f in image_files if f.lower().endswith(valid_exts)]
82
83 if not image_files:
84 return None
85 print(
86 f"Processing '{os.path.basename(folder_path)}' ({len(image_files)} images)..."
87 )
88
89 vectors = []
90 for img_path in image_files:
91 try:
92 img = Image.open(img_path).convert("RGB")
93 img_t = transform(img).unsqueeze(0).to(DEVICE)
94 with torch.no_grad():
95 output = model.forward_features(img_t)
96 vectors.append(output["x_norm_clstoken"].cpu())
97 except Exception:
98 continue
99
100 if not vectors:
101 return None
102
103 all_vecs = torch.cat(vectors, dim=0)
104 centroid = torch.mean(all_vecs, dim=0)
105 centroid = torch.nn.functional.normalize(centroid, p=2, dim=0)
106 return centroid
107
108
109class CentroidsWrapper(torch.nn.Module):
110 def __init__(self, data_dict):
111 super().__init__()
112 self.keys: List[str] = list(data_dict.keys())
113 self.tensors = torch.nn.ParameterList(
114 [torch.nn.Parameter(data_dict[k]) for k in self.keys]
115 )
116
117 def forward(self) -> Dict[str, torch.Tensor]:
118 result: Dict[str, torch.Tensor] = {}
119 for i, tensor in enumerate(self.tensors):
120 result[self.keys[i]] = tensor
121 return result
122
123
124# Runner
125
126if not os.path.exists(DATASET_ROOT):
127 print(f" Error: Dataset path not found: {DATASET_ROOT}")
128else:
129 centroids_dict = {}
130 folders = sorted(
131 [
132 d
133 for d in os.listdir(DATASET_ROOT)
134 if os.path.isdir(os.path.join(DATASET_ROOT, d))
135 ]
136 )
137
138 print(f"Found {len(folders)} texture categories.")
139
140 for category in folders:
141 full_path = os.path.join(DATASET_ROOT, category)
142 centroid = get_folder_centroid(full_path)
143 if centroid is not None:
144 centroids_dict[category] = centroid
145
146 if centroids_dict:
147 print(f"\nSaving {len(centroids_dict)} centroids to {OUTPUT_PATH}...")
148
149 wrapper = CentroidsWrapper(centroids_dict)
150 scripted_wrapper = torch.jit.script(wrapper)
151 scripted_wrapper.save(OUTPUT_PATH)
152
153 print(" SUCCESS! File saved.")
154 print(f"Output: {OUTPUT_PATH}")
155 else:
156 print(" Failed to generate any centroids.")
157
158
159## Texture Scoring
160import torch
161import torchvision.transforms as T
162import torch.nn.functional as F
163from PIL import Image
164import numpy as np
165import os
166
167DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
168PATCH_SIZE = 14
169CENTROIDS_PATH = "CENTROIDS_PATH"
170TEST_IMAGE_PATH = "TEST_IMAGE_PATH"
171
172
173class TextureScoreScanner:
174 def __init__(self, centroids_dict):
175 self.device = DEVICE
176 print(f"Loading DINOv2 (S14 - Small) on {self.device}...")
177 self.model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
178 self.model = self.model.half().to(self.device)
179 self.model.eval()
180
181 self.texture_names = []
182 matrix_list = []
183 print("Processing centroids...")
184 for name, tensor in centroids_dict.items():
185 if tensor.dim() > 1:
186 tensor = tensor.squeeze()
187 tensor = tensor.half().to(self.device)
188
189 self.texture_names.append(name)
190 matrix_list.append(tensor)
191
192 if len(matrix_list) > 0:
193 matrix = torch.stack(matrix_list)
194 self.centroid_matrix = F.normalize(matrix, p=2, dim=1)
195
196 if self.centroid_matrix.shape[1] != 384:
197 print(
198 f" WARNING: Dimension Mismatch! S14 expects 384, got {self.centroid_matrix.shape[1]}"
199 )
200 else:
201 raise ValueError("No valid centroids found in the .pt file.")
202
203 def preprocess(self, img_path):
204 if not os.path.exists(img_path):
205 raise FileNotFoundError(f"Image not found: {img_path}")
206 img = Image.open(img_path).convert("RGB")
207 w, h = img.size
208
209 new_w = (w // PATCH_SIZE) * PATCH_SIZE
210 new_h = (h // PATCH_SIZE) * PATCH_SIZE
211
212 transform = T.Compose(
213 [
214 T.Resize((new_h, new_w)),
215 T.ToTensor(),
216 T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
217 ]
218 )
219 return transform(img).unsqueeze(0).half().to(self.device)
220
221 def get_scores(self, img_path):
222 input_tensor = self.preprocess(img_path)
223
224 with torch.no_grad():
225 features = self.model.forward_features(input_tensor)
226
227 patches = features["x_norm_patchtokens"].squeeze(0)
228 patches_norm = F.normalize(patches, p=2, dim=1)
229
230 similarity_matrix = torch.mm(patches_norm, self.centroid_matrix.T)
231
232 class_scores = torch.mean(similarity_matrix, dim=0).cpu().numpy()
233
234 results = {}
235 for i, score in enumerate(class_scores):
236 results[self.texture_names[i]] = float(score)
237
238 return sorted(results.items(), key=lambda item: item[1], reverse=True)
239
240 def get_filtered_predictions(
241 self, sorted_scores, max_count=5, relative_thresh=0.85, absolute_floor=0.15
242 ):
243 if not sorted_scores:
244 return []
245
246 final_list = []
247 top_score = sorted_scores[0][1]
248
249 for name, score in sorted_scores:
250 if len(final_list) >= max_count:
251 break
252
253 if score < absolute_floor:
254 continue
255 if score < (top_score * relative_thresh):
256 break
257
258 final_list.append((name, score))
259
260 return final_list
261
262
263# Runner
264
265if os.path.exists(CENTROIDS_PATH):
266 print(f"Loading dictionary from {CENTROIDS_PATH}...")
267
268 loaded_container = torch.jit.load(CENTROIDS_PATH, map_location="cpu")
269 centroids_data = loaded_container()
270
271 scanner = TextureScoreScanner(centroids_data)
272
273 print(f"\nAnalyzing: {TEST_IMAGE_PATH}")
274 if os.path.exists(TEST_IMAGE_PATH):
275 raw_scores = scanner.get_scores(TEST_IMAGE_PATH)
276
277 smart_results = scanner.get_filtered_predictions(
278 raw_scores,
279 max_count=3, # Show max 3 textures
280 relative_thresh=0.85, # Must be 85% as good as the top match
281 absolute_floor=0.15, # Minimum similarity score
282 )
283
284 print("\n" + "=" * 40)
285 print(f"DETECTED TEXTURES ({len(smart_results)})")
286 print("=" * 40)
287
288 if len(smart_results) > 0:
289 top_score_val = raw_scores[0][1]
290 for name, score in smart_results:
291 percentage = (score / top_score_val) * 100
292 print(
293 f"• {name.upper().ljust(15)} : Score {score:.4f} (Conf: {percentage:.0f}%)"
294 )
295 else:
296 print("No significant texture match found.")
297 else:
298 print(f"Error: Test image not found at {TEST_IMAGE_PATH}")
299
300else:
301 print(f" Error: Centroids file '{CENTROIDS_PATH}' not found.")
302 print("Please run the training/generation script first.")
303
304## Convert DINOv2-S14 to FP16
305
306import torch
307import os
308
309MODEL_NAME = "dinov2_vits14"
310OUTPUT_FILENAME = "dinov2_vits14_fp16.pt"
311
312print(f" Downloading {MODEL_NAME} (FP32) from Torch Hub...")
313model = torch.hub.load("facebookresearch/dinov2", MODEL_NAME)
314
315print(" Converting model to Half Precision (FP16)...")
316model = model.half()
317
318print(f" Saving weights to {OUTPUT_FILENAME}...")
319torch.save(model.state_dict(), OUTPUT_FILENAME)
320
321# Verify
322file_size_mb = os.path.getsize(OUTPUT_FILENAME) / (1024 * 1024)
323print("-" * 30)
324print(f" Done! Saved to: {os.path.abspath(OUTPUT_FILENAME)}")
325print(f" File Size: {file_size_mb:.2f} MB")
326print("-" * 30)
327
328## Texture Scoring with FP16 model
329import torch
330import torchvision.transforms as T
331import torch.nn.functional as F
332from PIL import Image
333import numpy as np
334import os
335
336DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
337PATCH_SIZE = 14
338CENTROIDS_PATH = "centroids_path"
339TEST_IMAGE_PATH = "test_image_path"
340MODEL_WEIGHTS_PATH = "/kaggle/working/dinov2_vits14_fp16.pt"
341
342
343class TextureScoreScanner:
344 def __init__(self, centroids_dict):
345 self.device = DEVICE
346 print(f"Loading DINOv2 (S14 - Small) on {self.device}...")
347
348 self.model = torch.hub.load(
349 "facebookresearch/dinov2", "dinov2_vits14", pretrained=False
350 )
351
352 self.model = self.model.half().to(self.device)
353
354 if os.path.exists(MODEL_WEIGHTS_PATH):
355 print(f"Loading local weights from: {MODEL_WEIGHTS_PATH}")
356 state_dict = torch.load(MODEL_WEIGHTS_PATH, map_location=self.device)
357 self.model.load_state_dict(state_dict)
358 else:
359 raise FileNotFoundError(f"Weights file not found at: {MODEL_WEIGHTS_PATH}")
360
361 self.model.eval()
362
363 self.texture_names = []
364 matrix_list = []
365
366 print("Processing centroids...")
367 for name, tensor in centroids_dict.items():
368 if tensor.dim() > 1:
369 tensor = tensor.squeeze()
370
371 tensor = tensor.half().to(self.device)
372
373 self.texture_names.append(name)
374 matrix_list.append(tensor)
375
376 if len(matrix_list) > 0:
377 matrix = torch.stack(matrix_list)
378 self.centroid_matrix = F.normalize(matrix, p=2, dim=1)
379
380 if self.centroid_matrix.shape[1] != 384:
381 print(
382 f"⚠️ WARNING: Dimension Mismatch! S14 expects 384, got {self.centroid_matrix.shape[1]}"
383 )
384 else:
385 raise ValueError("No valid centroids found in the .pt file.")
386
387 def preprocess(self, img_path):
388 if not os.path.exists(img_path):
389 raise FileNotFoundError(f"Image not found: {img_path}")
390 img = Image.open(img_path).convert("RGB")
391 w, h = img.size
392
393 new_w = (w // PATCH_SIZE) * PATCH_SIZE
394 new_h = (h // PATCH_SIZE) * PATCH_SIZE
395
396 transform = T.Compose(
397 [
398 T.Resize((new_h, new_w)),
399 T.ToTensor(),
400 T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
401 ]
402 )
403
404 return transform(img).unsqueeze(0).half().to(self.device)
405
406 def get_scores(self, img_path):
407 input_tensor = self.preprocess(img_path)
408
409 with torch.no_grad():
410 features = self.model.forward_features(input_tensor)
411
412 patches = features["x_norm_patchtokens"].squeeze(0)
413 patches_norm = F.normalize(patches, p=2, dim=1)
414
415 similarity_matrix = torch.mm(patches_norm, self.centroid_matrix.T)
416
417 class_scores = torch.mean(similarity_matrix, dim=0).float().cpu().numpy()
418
419 results = {}
420 for i, score in enumerate(class_scores):
421 results[self.texture_names[i]] = float(score)
422
423 return sorted(results.items(), key=lambda item: item[1], reverse=True)
424
425 def get_filtered_predictions(
426 self, sorted_scores, max_count=5, relative_thresh=0.85, absolute_floor=0.15
427 ):
428 if not sorted_scores:
429 return []
430
431 final_list = []
432 top_score = sorted_scores[0][1]
433
434 for name, score in sorted_scores:
435 if len(final_list) >= max_count:
436 break
437
438 if score < absolute_floor:
439 continue
440 if score < (top_score * relative_thresh):
441 break
442
443 final_list.append((name, score))
444
445 return final_list
446
447
448# Runner
449
450if os.path.exists(CENTROIDS_PATH):
451 print(f"Loading dictionary from {CENTROIDS_PATH}...")
452
453 loaded_container = torch.jit.load(CENTROIDS_PATH, map_location="cpu")
454 centroids_data = loaded_container()
455
456 scanner = TextureScoreScanner(centroids_data)
457
458 print(f"\nAnalyzing: {TEST_IMAGE_PATH}")
459 if os.path.exists(TEST_IMAGE_PATH):
460 raw_scores = scanner.get_scores(TEST_IMAGE_PATH)
461
462 smart_results = scanner.get_filtered_predictions(
463 raw_scores,
464 max_count=3,
465 relative_thresh=0.85,
466 absolute_floor=0.05, # Kept at 0.05 as discussed
467 )
468
469 print("\n" + "=" * 40)
470 print(f"DETECTED TEXTURES ({len(smart_results)})")
471 print("=" * 40)
472
473 if len(smart_results) > 0:
474 top_score_val = raw_scores[0][1]
475 for name, score in smart_results:
476 percentage = (score / top_score_val) * 100
477 print(
478 f"• {name.upper().ljust(15)} : Score {score:.4f} (Conf: {percentage:.0f}%)"
479 )
480 else:
481 print("No significant texture match found.")
482 else:
483 print(f"Error: Test image not found at {TEST_IMAGE_PATH}")
484
485else:
486 print(f" Error: Centroids file '{CENTROIDS_PATH}' not found.")