load image with resnet
This commit is contained in:
27
train/resnet01.py
Normal file
27
train/resnet01.py
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
import torch
|
||||||
|
import torchvision.models as models
|
||||||
|
from PIL import Image
|
||||||
|
import torchvision.transforms as transforms
|
||||||
|
|
||||||
|
model = models.resnet50(pretrained=True)
|
||||||
|
|
||||||
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||||
|
model = model.to(device)
|
||||||
|
|
||||||
|
image_path = '1706678996799-2.jpeg'
|
||||||
|
image = Image.open(image_path).convert('RGB')
|
||||||
|
|
||||||
|
transform = transforms.Compose([
|
||||||
|
transforms.Resize((224, 224)),
|
||||||
|
transforms.ToTensor(),
|
||||||
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||||
|
])
|
||||||
|
|
||||||
|
input_data = transform(image).unsqueeze(0).to(device)
|
||||||
|
|
||||||
|
model.eval()
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
output = model(input_data)
|
||||||
|
|
||||||
|
print(output)
|
||||||
Reference in New Issue
Block a user