|
|
|
|
|
import torch |
|
import torchvision |
|
from torchvision.models import resnet50 |
|
from torchvision.transforms import transforms |
|
from torch.utils.data import DataLoader, Dataset |
|
from PIL import Image |
|
|
|
class RobloxDataset(Dataset): |
|
def __init__(self, root_dir, transform=None): |
|
self.root_dir = root_dir |
|
self.transform = transform |
|
|
|
def __len__(self): |
|
return 15 |
|
|
|
def __getitem__(self, idx): |
|
img_path = f'{self.root_dir}/Character/human{str(idx+1).zfill(2)}.png' |
|
image = Image.open(img_path).convert('RGB') |
|
|
|
if self.transform: |
|
image = self.transform(image) |
|
|
|
return image |
|
|
|
data_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]) |
|
]) |
|
|
|
dataset = RobloxDataset('/content/dataset', transform=data_transform) |
|
|
|
data_loader = DataLoader(dataset, batch_size=1, shuffle=True) |
|
|
|
model = resnet50(pretrained=True) |
|
model.fc = torch.nn.Linear(in_features=2048, out_features=1) |
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
model.to(device) |
|
|
|
criterion = torch.nn.BCEWithLogitsLoss() |
|
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) |
|
|
|
num_epochs = 10 |
|
|
|
for epoch in range(num_epochs): |
|
for images in data_loader: |
|
images = images.to(device) |
|
labels = torch.ones((images.size(0), 1)).to(device) |
|
|
|
outputs = model(images) |
|
loss = criterion(outputs, labels) |
|
|
|
optimizer.zero_grad() |
|
loss.backward() |
|
optimizer.step() |
|
|
|
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}') |
