Spaces:
Sleeping
Sleeping
File size: 870 Bytes
ab9b7a8 |
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 |
import os
from tqdm.auto import tqdm
from PIL import Image
import torch
import torch.nn as nn
from torchvision import models, transforms
device = "cuda" if torch.cuda.is_available() else "cpu"
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor()
])
def preprocess_image(image_path):
img = Image.open(image_path).convert('RGB')
processed_img = transform(img)
return processed_img
def create_resnet18_model():
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
modules = list(model.children())[:-1]
model = nn.Sequential(*modules)
return model
def extract_features(model, processed_image):
input = processed_image.unsqueeze(dim=0).to(device)
model.eval()
with torch.inference_mode():
prediction = model(input)
return prediction.squeeze().tolist()
|