diff --git a/train/resnet01.py b/train/resnet01.py new file mode 100644 index 0000000..d1a835a --- /dev/null +++ b/train/resnet01.py @@ -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) \ No newline at end of file