creek

# -*- coding: utf-8 -*-
"""
Texture and Material
Original file: https://colab.research.google.com/drive/1u-5cHHckvZPMjIFJ5YJcCAGzDWVPRl91
"""

## .pth -> .pt

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import numpy as np
import os

CENTROIDS_PATH = "texture_centroids_augmented.pth_path"
OUTPUT_PATH = "texture_centroids_augmented.pt_path"


class DataWrapper(nn.Module):
    def __init__(self, data_dict):
        super().__init__()
        self.data = torch.jit.Attribute(data_dict, dict[str, torch.Tensor])

    def forward(self):
        return self.data


if os.path.exists(CENTROIDS_PATH):
    print("Loading .pth file...")
    raw_data = torch.load(CENTROIDS_PATH, map_location="cpu", weights_only=False)

    print("Processing data...")
    clean_dict = {}
    for k, v in raw_data.items():
        if isinstance(v, np.ndarray):
            clean_dict[k] = torch.from_numpy(v).float()
        else:
            clean_dict[k] = v.float()

    wrapper = DataWrapper(clean_dict)
    scripted_wrapper = torch.jit.script(wrapper)
    scripted_wrapper.save(OUTPUT_PATH)
    print(f"Success! Saved as {OUTPUT_PATH}")
else:
    print("File not found.")

## Generate centroids

import torch
import torchvision.transforms as T
from PIL import Image
import numpy as np
import os
import glob
from typing import Dict, List

DATASET_ROOT = "Texture_Dataset_path"
OUTPUT_PATH = "texture_centroids_s14.pt"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


print(f"Loading DINOv2-S14 on {DEVICE}...")
model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
model.to(DEVICE)
model.eval()

transform = T.Compose(
    [
        T.Resize(256, interpolation=T.InterpolationMode.BICUBIC),
        T.CenterCrop(224),
        T.ToTensor(),
        T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ]
)


def get_folder_centroid(folder_path):
    image_files = glob.glob(os.path.join(folder_path, "*"))
    valid_exts = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
    image_files = [f for f in image_files if f.lower().endswith(valid_exts)]

    if not image_files:
        return None
    print(
        f"Processing '{os.path.basename(folder_path)}' ({len(image_files)} images)..."
    )

    vectors = []
    for img_path in image_files:
        try:
            img = Image.open(img_path).convert("RGB")
            img_t = transform(img).unsqueeze(0).to(DEVICE)
            with torch.no_grad():
                output = model.forward_features(img_t)
                vectors.append(output["x_norm_clstoken"].cpu())
        except Exception:
            continue

    if not vectors:
        return None

    all_vecs = torch.cat(vectors, dim=0)
    centroid = torch.mean(all_vecs, dim=0)
    centroid = torch.nn.functional.normalize(centroid, p=2, dim=0)
    return centroid


class CentroidsWrapper(torch.nn.Module):
    def __init__(self, data_dict):
        super().__init__()
        self.keys: List[str] = list(data_dict.keys())
        self.tensors = torch.nn.ParameterList(
            [torch.nn.Parameter(data_dict[k]) for k in self.keys]
        )

    def forward(self) -> Dict[str, torch.Tensor]:
        result: Dict[str, torch.Tensor] = {}
        for i, tensor in enumerate(self.tensors):
            result[self.keys[i]] = tensor
        return result


#  Runner

if not os.path.exists(DATASET_ROOT):
    print(f" Error: Dataset path not found: {DATASET_ROOT}")
else:
    centroids_dict = {}
    folders = sorted(
        [
            d
            for d in os.listdir(DATASET_ROOT)
            if os.path.isdir(os.path.join(DATASET_ROOT, d))
        ]
    )

    print(f"Found {len(folders)} texture categories.")

    for category in folders:
        full_path = os.path.join(DATASET_ROOT, category)
        centroid = get_folder_centroid(full_path)
        if centroid is not None:
            centroids_dict[category] = centroid

    if centroids_dict:
        print(f"\nSaving {len(centroids_dict)} centroids to {OUTPUT_PATH}...")

        wrapper = CentroidsWrapper(centroids_dict)
        scripted_wrapper = torch.jit.script(wrapper)
        scripted_wrapper.save(OUTPUT_PATH)

        print(" SUCCESS! File saved.")
        print(f"Output: {OUTPUT_PATH}")
    else:
        print(" Failed to generate any centroids.")


## Texture Scoring
import torch
import torchvision.transforms as T
import torch.nn.functional as F
from PIL import Image
import numpy as np
import os

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
PATCH_SIZE = 14
CENTROIDS_PATH = "CENTROIDS_PATH"
TEST_IMAGE_PATH = "TEST_IMAGE_PATH"


class TextureScoreScanner:
    def __init__(self, centroids_dict):
        self.device = DEVICE
        print(f"Loading DINOv2 (S14 - Small) on {self.device}...")
        self.model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
        self.model = self.model.half().to(self.device)
        self.model.eval()

        self.texture_names = []
        matrix_list = []
        print("Processing centroids...")
        for name, tensor in centroids_dict.items():
            if tensor.dim() > 1:
                tensor = tensor.squeeze()
            tensor = tensor.half().to(self.device)

            self.texture_names.append(name)
            matrix_list.append(tensor)

        if len(matrix_list) > 0:
            matrix = torch.stack(matrix_list)
            self.centroid_matrix = F.normalize(matrix, p=2, dim=1)

            if self.centroid_matrix.shape[1] != 384:
                print(
                    f" WARNING: Dimension Mismatch! S14 expects 384, got {self.centroid_matrix.shape[1]}"
                )
        else:
            raise ValueError("No valid centroids found in the .pt file.")

    def preprocess(self, img_path):
        if not os.path.exists(img_path):
            raise FileNotFoundError(f"Image not found: {img_path}")
        img = Image.open(img_path).convert("RGB")
        w, h = img.size

        new_w = (w // PATCH_SIZE) * PATCH_SIZE
        new_h = (h // PATCH_SIZE) * PATCH_SIZE

        transform = T.Compose(
            [
                T.Resize((new_h, new_w)),
                T.ToTensor(),
                T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            ]
        )
        return transform(img).unsqueeze(0).half().to(self.device)

    def get_scores(self, img_path):
        input_tensor = self.preprocess(img_path)

        with torch.no_grad():
            features = self.model.forward_features(input_tensor)

            patches = features["x_norm_patchtokens"].squeeze(0)
            patches_norm = F.normalize(patches, p=2, dim=1)

            similarity_matrix = torch.mm(patches_norm, self.centroid_matrix.T)

            class_scores = torch.mean(similarity_matrix, dim=0).cpu().numpy()

            results = {}
            for i, score in enumerate(class_scores):
                results[self.texture_names[i]] = float(score)

            return sorted(results.items(), key=lambda item: item[1], reverse=True)

    def get_filtered_predictions(
        self, sorted_scores, max_count=5, relative_thresh=0.85, absolute_floor=0.15
    ):
        if not sorted_scores:
            return []

        final_list = []
        top_score = sorted_scores[0][1]

        for name, score in sorted_scores:
            if len(final_list) >= max_count:
                break

            if score < absolute_floor:
                continue
            if score < (top_score * relative_thresh):
                break

            final_list.append((name, score))

        return final_list


# Runner

if os.path.exists(CENTROIDS_PATH):
    print(f"Loading dictionary from {CENTROIDS_PATH}...")

    loaded_container = torch.jit.load(CENTROIDS_PATH, map_location="cpu")
    centroids_data = loaded_container()

    scanner = TextureScoreScanner(centroids_data)

    print(f"\nAnalyzing: {TEST_IMAGE_PATH}")
    if os.path.exists(TEST_IMAGE_PATH):
        raw_scores = scanner.get_scores(TEST_IMAGE_PATH)

        smart_results = scanner.get_filtered_predictions(
            raw_scores,
            max_count=3,  # Show max 3 textures
            relative_thresh=0.85,  # Must be 85% as good as the top match
            absolute_floor=0.15,  # Minimum similarity score
        )

        print("\n" + "=" * 40)
        print(f"DETECTED TEXTURES ({len(smart_results)})")
        print("=" * 40)

        if len(smart_results) > 0:
            top_score_val = raw_scores[0][1]
            for name, score in smart_results:
                percentage = (score / top_score_val) * 100
                print(
                    f"• {name.upper().ljust(15)} : Score {score:.4f} (Conf: {percentage:.0f}%)"
                )
        else:
            print("No significant texture match found.")
    else:
        print(f"Error: Test image not found at {TEST_IMAGE_PATH}")

else:
    print(f" Error: Centroids file '{CENTROIDS_PATH}' not found.")
    print("Please run the training/generation script first.")

## Convert DINOv2-S14 to FP16

import torch
import os

MODEL_NAME = "dinov2_vits14"
OUTPUT_FILENAME = "dinov2_vits14_fp16.pt"

print(f"  Downloading {MODEL_NAME} (FP32) from Torch Hub...")
model = torch.hub.load("facebookresearch/dinov2", MODEL_NAME)

print(" Converting model to Half Precision (FP16)...")
model = model.half()

print(f" Saving weights to {OUTPUT_FILENAME}...")
torch.save(model.state_dict(), OUTPUT_FILENAME)

# Verify
file_size_mb = os.path.getsize(OUTPUT_FILENAME) / (1024 * 1024)
print("-" * 30)
print(f" Done! Saved to: {os.path.abspath(OUTPUT_FILENAME)}")
print(f" File Size: {file_size_mb:.2f} MB")
print("-" * 30)

## Texture Scoring with FP16 model
import torch
import torchvision.transforms as T
import torch.nn.functional as F
from PIL import Image
import numpy as np
import os

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
PATCH_SIZE = 14
CENTROIDS_PATH = "centroids_path"
TEST_IMAGE_PATH = "test_image_path"
MODEL_WEIGHTS_PATH = "/kaggle/working/dinov2_vits14_fp16.pt"


class TextureScoreScanner:
    def __init__(self, centroids_dict):
        self.device = DEVICE
        print(f"Loading DINOv2 (S14 - Small) on {self.device}...")

        self.model = torch.hub.load(
            "facebookresearch/dinov2", "dinov2_vits14", pretrained=False
        )

        self.model = self.model.half().to(self.device)

        if os.path.exists(MODEL_WEIGHTS_PATH):
            print(f"Loading local weights from: {MODEL_WEIGHTS_PATH}")
            state_dict = torch.load(MODEL_WEIGHTS_PATH, map_location=self.device)
            self.model.load_state_dict(state_dict)
        else:
            raise FileNotFoundError(f"Weights file not found at: {MODEL_WEIGHTS_PATH}")

        self.model.eval()

        self.texture_names = []
        matrix_list = []

        print("Processing centroids...")
        for name, tensor in centroids_dict.items():
            if tensor.dim() > 1:
                tensor = tensor.squeeze()

            tensor = tensor.half().to(self.device)

            self.texture_names.append(name)
            matrix_list.append(tensor)

        if len(matrix_list) > 0:
            matrix = torch.stack(matrix_list)
            self.centroid_matrix = F.normalize(matrix, p=2, dim=1)

            if self.centroid_matrix.shape[1] != 384:
                print(
                    f"⚠️ WARNING: Dimension Mismatch! S14 expects 384, got {self.centroid_matrix.shape[1]}"
                )
        else:
            raise ValueError("No valid centroids found in the .pt file.")

    def preprocess(self, img_path):
        if not os.path.exists(img_path):
            raise FileNotFoundError(f"Image not found: {img_path}")
        img = Image.open(img_path).convert("RGB")
        w, h = img.size

        new_w = (w // PATCH_SIZE) * PATCH_SIZE
        new_h = (h // PATCH_SIZE) * PATCH_SIZE

        transform = T.Compose(
            [
                T.Resize((new_h, new_w)),
                T.ToTensor(),
                T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            ]
        )

        return transform(img).unsqueeze(0).half().to(self.device)

    def get_scores(self, img_path):
        input_tensor = self.preprocess(img_path)

        with torch.no_grad():
            features = self.model.forward_features(input_tensor)

            patches = features["x_norm_patchtokens"].squeeze(0)
            patches_norm = F.normalize(patches, p=2, dim=1)

            similarity_matrix = torch.mm(patches_norm, self.centroid_matrix.T)

            class_scores = torch.mean(similarity_matrix, dim=0).float().cpu().numpy()

            results = {}
            for i, score in enumerate(class_scores):
                results[self.texture_names[i]] = float(score)

            return sorted(results.items(), key=lambda item: item[1], reverse=True)

    def get_filtered_predictions(
        self, sorted_scores, max_count=5, relative_thresh=0.85, absolute_floor=0.15
    ):
        if not sorted_scores:
            return []

        final_list = []
        top_score = sorted_scores[0][1]

        for name, score in sorted_scores:
            if len(final_list) >= max_count:
                break

            if score < absolute_floor:
                continue
            if score < (top_score * relative_thresh):
                break

            final_list.append((name, score))

        return final_list


# Runner

if os.path.exists(CENTROIDS_PATH):
    print(f"Loading dictionary from {CENTROIDS_PATH}...")

    loaded_container = torch.jit.load(CENTROIDS_PATH, map_location="cpu")
    centroids_data = loaded_container()

    scanner = TextureScoreScanner(centroids_data)

    print(f"\nAnalyzing: {TEST_IMAGE_PATH}")
    if os.path.exists(TEST_IMAGE_PATH):
        raw_scores = scanner.get_scores(TEST_IMAGE_PATH)

        smart_results = scanner.get_filtered_predictions(
            raw_scores,
            max_count=3,
            relative_thresh=0.85,
            absolute_floor=0.05,  # Kept at 0.05 as discussed
        )

        print("\n" + "=" * 40)
        print(f"DETECTED TEXTURES ({len(smart_results)})")
        print("=" * 40)

        if len(smart_results) > 0:
            top_score_val = raw_scores[0][1]
            for name, score in smart_results:
                percentage = (score / top_score_val) * 100
                print(
                    f"• {name.upper().ljust(15)} : Score {score:.4f} (Conf: {percentage:.0f}%)"
                )
        else:
            print("No significant texture match found.")
    else:
        print(f"Error: Test image not found at {TEST_IMAGE_PATH}")

else:
    print(f" Error: Centroids file '{CENTROIDS_PATH}' not found.")