model.py

import torch
import torchvision.transforms as T
from PIL import Image
import joblib
import json
import cv2
import gradio as gr

# Define image transformation
transform_image = T.Compose([
    T.ToTensor(),
    T.Resize(244),
    T.CenterCrop(224),
    T.Normalize([0.5], [0.5])
])

def load_image(img: str) -> torch.Tensor:
    """
    Load an image and return a tensor that can be used as an input to DINOv2.
    """
    img = Image.open(img)
    transformed_img = transform_image(img)[:3].unsqueeze(0)
    return transformed_img

# Load models for inference
dinov2_vits14 = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
dinov2_vits14.to(device)
dinov2_vits14.eval()  # Set the model to evaluation mode

# Load the classifier
clf = joblib.load('svm_model.joblib')

# Load the embeddings
with open('all_embeddings.json', 'r') as f:
    embeddings = json.load(f)

# Predict class for a new image
def predict(image_path):
    new_image = load_image(image_path).to(device)
    with torch.no_grad():
        embedding = dinov2_vits14(new_image).cpu().numpy().reshape(1, -1)
    prediction = clf.predict(embedding)
    return prediction[0]