Spaces:
Runtime error
Runtime error
File size: 2,633 Bytes
54660f7 |
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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision.models as models
import torchvision.transforms.functional as VF
from torchvision import transforms
import sys, argparse, os, glob
import pandas as pd
import numpy as np
from PIL import Image
from collections import OrderedDict
class ToPIL(object):
def __call__(self, sample):
img = sample
img = transforms.functional.to_pil_image(img)
return img
class BagDataset():
def __init__(self, csv_file, transform=None):
self.files_list = csv_file
self.transform = transform
def __len__(self):
return len(self.files_list)
def __getitem__(self, idx):
temp_path = self.files_list[idx]
img = os.path.join(temp_path)
img = Image.open(img)
img = img.resize((224, 224))
sample = {'input': img}
if self.transform:
sample = self.transform(sample)
return sample
class ToTensor(object):
def __call__(self, sample):
img = sample['input']
img = VF.to_tensor(img)
return {'input': img}
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, img):
for t in self.transforms:
img = t(img)
return img
def save_coords(txt_file, csv_file_path):
for path in csv_file_path:
x, y = path.split('/')[-1].split('.')[0].split('_')
txt_file.writelines(str(x) + '\t' + str(y) + '\n')
txt_file.close()
def adj_matrix(csv_file_path, output, device='cpu'):
total = len(csv_file_path)
adj_s = np.zeros((total, total))
for i in range(total-1):
path_i = csv_file_path[i]
x_i, y_i = path_i.split('/')[-1].split('.')[0].split('_')
for j in range(i+1, total):
# sptial
path_j = csv_file_path[j]
x_j, y_j = path_j.split('/')[-1].split('.')[0].split('_')
if abs(int(x_i)-int(x_j)) <=1 and abs(int(y_i)-int(y_j)) <= 1:
adj_s[i][j] = 1
adj_s[j][i] = 1
adj_s = torch.from_numpy(adj_s)
adj_s = adj_s.to(device)
return adj_s
def bag_dataset(args, csv_file_path):
transformed_dataset = BagDataset(csv_file=csv_file_path,
transform=Compose([
ToTensor()
]))
dataloader = DataLoader(transformed_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, drop_last=False)
return dataloader, len(transformed_dataset) |