import os
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from transformers import ConvNextImageProcessor, ConvNextForImageClassification

#Classe para extração de features (ConvNeXt)
class FeatureExtractor:
    def __init__(self, device=None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Usando dispositivo: {self.device}")

        #Modelo e pré-processador
        self.processor = ConvNextImageProcessor.from_pretrained(
            "facebook/convnext-large-224-22k-1k"
        )
        self.model = ConvNextForImageClassification.from_pretrained(
            "facebook/convnext-large-224-22k-1k"
        ).to(self.device)

        #Remove camada de classificação (ficam só as features)
        self.model.classifier = torch.nn.Identity()
        self.model.eval()

    def extract_convnext(self, image_path: str) -> np.ndarray:
        image = Image.open(image_path).convert("RGB")
        inputs = self.processor(image, return_tensors="pt").to(self.device)
        with torch.no_grad():
            features = self.model(**inputs).logits
        features_np = features.cpu().numpy().flatten()
        print(f"Vetor de características extraído com shape: {features_np.shape}")
        return features_np

#Função principal
def process_single_image(image_path: str, output_dir: str = "processed"):
    """
    Extração de features ConvNeXt (sem segmentação)
    """
    #1) Extração de features
    extractor = FeatureExtractor()
    #Passa o CAMINHO ORIGINAL da imagem (ex: /tmp/temp_foto.jpg)
    features = extractor.extract_convnext(image_path) 

    #NÃO PRECISAMOS MAIS SALVAR O ARQUIVO .NPY
    #feat_path = os.path.join(output_dir, "features_single.npy")
    #np.save(feat_path, features)
    #print(f"Vetor de características salvo em: {feat_path}")

    return features #RETORNA O ARRAY DE DADOS

#Execução direta
if __name__ == "__main__":
    #Caminho da imagem
    image_path = r"C:\Users\Julia\Downloads\clb2_600px.jpg"

    #Executa pipeline completo
    process_single_image(image_path)