def main():
fixed_seed(42)
if args.mode == 'train':
model = get_model_instance_segmentation(args.num_classes).to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.SGD(params, lr=1e-3, momentum=0.9, weight_decay=1e-4)
for epoch in range(args.epochs):
train_fn(model, train_loader, optimizer, device, epoch)
torch.save(model.state_dict(), f'./weight/test{epoch}.pth')
else:
model = get_model_instance_segmentation(args.num_classes).to(device)
model.load_state_dict(torch.load('./weight/test1.pth'))
model.eval()
test_fn(model, test_loader, class_nums, device)
def main():
fixed_seed(42)
model = get_model_instance_segmentation(args.num_classes).to(device)
model.load_state_dict(torch.load('./weight/test1.pth'))
model.eval()
test_fn(model, test_loader, class_nums, device)
def predict(image):
model = get_model_instance_segmentation(2)
model.load_state_dict(
torch.load('best_model', map_location=torch.device('cpu')))
model.eval()
image = Image.open(image).convert('RGB')
img = loader(image)
with torch.no_grad():
output = model(img[None, ...])[0]
display_instances(image, output['boxes'], output['labels'], class_names,
output['scores'])
def run_train():
fixed_seed(42)
if not os.path.exists('./weight/'):
print('Not exists ./weight/ making an weight folder...')
os.mkdir('./weight/')
model = get_model_instance_segmentation(args.num_classes).to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.SGD(params, lr=1e-3, momentum=0.9, weight_decay=1e-4)
for epoch in range(args.epochs):
train_fn(model, train_loader, optimizer, device, epoch)
torch.save(model.state_dict(), f'./weight/test{epoch}.pth')
def main():
parser = argparse.ArgumentParser(description='Mask R-CNN')
parser.add_argument(
'--image_path',
default=
'/home/hatsunemiku/dev/mask-rcnn/data/food201/segmented_test/pizza/309892.jpg',
type=str,
help='image path')
args = parser.parse_args()
model = get_model_instance_segmentation(num_classes=2)
model.load_state_dict(torch.load('food201_model.pkl'))
model.eval()
detect_and_color_splash(model, args.image_path, dataset_name='food201')
def train(config):
# create our own dataset and its data_loader
tr_dt = FashionDataset(config, get_transform(train = True))
tr_data_loader = DataLoader(
tr_dt, config.batch_size , shuffle = True,
num_workers = 8, collate_fn = lambda x: tuple(zip(*x))
)
# save the weights
# weight_file = osp.join(config.save_dir, 'weights')
# check them whether exists or not
if not os.path.exists(config.save_dir):
os.makedirs(config.save_dir)
# there are 46 classes in total
num_classes = 46 + 1
# create model instance
model = get_model_instance_segmentation(num_classes)
#set model to device
model.to(device)
# for optim
params = [p for p in model.parameters() if p.requires_grad]
optim = torch.optim.SGD(params, lr = 0.001, momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optim,
step_size=5,
gamma=0.1)
for epoch in range(config.num_epochs):
train_one_epoch(model, optim, tr_data_loader, device, epoch, print_freq=config.rep_intv)
# updt the learning rate
lr_scheduler.step()
w1 = osp.join(config.save_dir , 'weights')
wfile = osp.join(w1, '{}_model.bin'.format(str(epoch)))
torch.save(model.state_dict(), wfile)
def main():
# train on the GPU or on the CPU, if a GPU is not available
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
print(f'Device: {device}')
# our dataset has two classes only - background and person
num_classes = 2
# use our dataset and defined transformations
dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))
# split the dataset in train and test set
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=1,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(dataset_test,
batch_size=1,
shuffle=False,
num_workers=1,
collate_fn=utils.collate_fn)
# get the model using our helper function
model = get_model_instance_segmentation(num_classes)
# move model to the right device
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
# let's train it for 10 epochs
num_epochs = 10
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model,
optimizer,
data_loader,
device,
epoch,
print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, data_loader_test, device=device)
print("That's it!")
# split the dataset in train and test set
#indices = torch.randperm(len(dataset)).tolist()
#dataset = torch.utils.data.Subset(dataset, indices)
#indices_test = torch.randperm(len(dataset_test)).tolist()
#dataset_test = torch.utils.data.Subset(dataset_test, indices_test)
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=4, shuffle=True, num_workers=8,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1, shuffle=False, num_workers=4,
collate_fn=utils.collate_fn)
# get the model using our helper function
mask_rcnn = get_model_instance_segmentation(num_classes, image_mean, image_std, stats=True)
read_param = False
if read_param:
mask_rcnn.load_state_dict(torch.load("trained_param_6/epoch_0050.param"))
print("Loaded weights")
# move model to the right device
mask_rcnn.to(device)
# construct an optimizer
params = [p for p in mask_rcnn.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=40,
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-b", "--backbone", type=str,
choices=['resnet', 'mn2'], default='resnet',
help="backbone type")
parser.add_argument("-m", "--modelpath",
default='./checkpoint/inb_resnet_9',
help="model file path")
parser.add_argument("-d", "--datapath", default='./benchmark',
help="data path")
args = parser.parse_args()
dataset = INBDataset(args.datapath, M.get_transform(train=False))
dataset_test = dataset
num_classes = 2
model = M.get_model_instance_segmentation(args.backbone, num_classes)
checkpoint = torch.load(args.modelpath)
# model.load_state_dict(checkpoint['model_state_dict'])
model.load_state_dict(checkpoint)
device = torch.device('cuda') if torch.cuda.is_available()\
else torch.device('cpu')
model.to(device)
stat(model, dataset_test, device)
batch_size=2,
shuffle=True,
num_workers=2,
collate_fn=dataset.collate_fn,
)
dataloader = {"train": dataloader_train, "val": dataloader_test}
if checkpoint is None:
# train on the GPU or on the CPU, if a GPU is not available
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
device = (torch.device("cuda:0")
if torch.cuda.is_available() else torch.device("cpu"))
print(device)
with torch.cuda.device(0):
# load a pre-trained model
model = get_model_instance_segmentation(21)
optimizer_ft = torch.optim.SGD(model.parameters(),
lr=0.01,
momentum=0.9,
weight_decay=0.0005)
# and a learning rate scheduler
milestones = [10, 23, 33, 43]
exp_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer_ft,
milestones,
gamma=0.1,
last_epoch=-1)
# exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_ft,
# step_size=10,
# gamma=0.1)
# our dataset has three classes only - background, non-damaged, and damaged
num_classes = 6 # 3 or 6
dataset_test = Dataset("./datasets/Eureka_infer/102/",
"./datasets/Eureka_infer/102_labels/",
get_transform(train=False),
readsave=False)
data_loader_test = torch.utils.data.DataLoader(dataset_test,
batch_size=1,
shuffle=False,
num_workers=2,
collate_fn=utils.collate_fn)
mask_rcnn = get_model_instance_segmentation(num_classes,
image_mean=None,
image_std=None,
stats=False)
mask_rcnn.load_state_dict(
torch.load("trained_param_eureka_aug_mult/epoch_0021.param"))
print("loaded weights")
# move model to the right device
mask_rcnn.to(device)
mask_rcnn.eval()
for image, target in dataset_test:
pred = mask_rcnn(image.unsqueeze(0).to(device))[0]
boxes = pred['boxes'].to("cpu").detach().numpy()
def main():
# train on the GPU or on the CPU, if a GPU is not available
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(device)
# our dataset has two classes only - background and person
num_classes = 6
# use our dataset and defined transformations
MASK_MAP = {"title": 1,
"rating": 2,
"ratings": 3,
"price": 4,
"description": 5}
dataset = LabelboxWebpageDataset("../data/product_images", "../data/labelbox-export.862Z.json", MASK_MAP,
get_transform(True))
dataset_test = LabelboxWebpageDataset("../data/product_images", "../data/labelbox-export.862Z.json", MASK_MAP,
get_transform(False))
# split the dataset in train and test set
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-5])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-5:])
# define training and validatio n data loaders
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=1, shuffle=True, num_workers=0,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1, shuffle=False, num_workers=0,
collate_fn=utils.collate_fn)
# get the model using our helper function
model = get_model_instance_segmentation(num_classes)
# move model to the right device
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
# optimizer = torch.optim.SGD(params, lr=0.005,
# momentum=0.9, weight_decay=0.0005)
optimizer = torch.optim.Adam(params, lr=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
# let's train it for 10 epochs
num_epochs = 10
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
# Evaluation broken
if epoch == num_epochs - 1:
model.eval()
# cpu_device = torch.device("cpu")
# model.to(cpu_device)
# for images, targets in data_loader_test:
# # images = list(img.to(cpu_device) for img in images)
# # outputs = model(images)
# # print(outputs)
# # print(targets)
evaluate(model, data_loader_test, device=device)
def main():
# train on the GPU or on the CPU, if a GPU is not available
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
# our dataset has two classes only - background and person
num_classes = 2
# get the model using our helper function
model = get_model_instance_segmentation(num_classes)
# load pretrain_dict
pretrain_dict = torch.load(
os.path.join("C:\\zhulei\\maskRcnn\\models", "_epoch-9.pth"))
model.load_state_dict(pretrain_dict)
# move model to the right device
model.to(device)
# use our dataset and defined transformations
dataset_test = PennFudanDataset('C:\\zhulei\\maskRcnn\\data\\test',
get_transform(train=False))
data_loader_test = torch.utils.data.DataLoader(dataset_test,
batch_size=1,
shuffle=False,
num_workers=1,
collate_fn=utils.collate_fn)
n_threads = torch.get_num_threads()
# FIXME remove this and make paste_masks_in_image run on the GPU
torch.set_num_threads(1)
cpu_device = torch.device("cpu")
model.eval()
metric_logger = utils.MetricLogger(delimiter=" ")
header = 'Test:'
coco = get_coco_api_from_dataset(data_loader_test.dataset)
iou_types = _get_iou_types(model)
coco_evaluator = CocoEvaluator(coco, iou_types)
for image, targets in metric_logger.log_every(data_loader_test, 100,
header):
image = list(img.to(device) for img in image)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
torch.cuda.synchronize()
model_time = time.time()
outputs = model(image)
instance_segmentation_api(image[0], outputs)
# 可视化
# for img in image:
# Image.fromarray((img.mul(255).permute(1, 2, 0).byte().cpu().numpy())[0])
# print(outputs[0]['masks'].shape)
# for i in range(99):
# result = Image.fromarray(outputs[0]['masks'][i, 0].mul(255).byte().cpu().numpy())
# result.show()
outputs = [{k: v.to(cpu_device)
for k, v in t.items()} for t in outputs]
model_time = time.time() - model_time
res = {
target["image_id"].item(): output
for target, output in zip(targets, outputs)
}
evaluator_time = time.time()
coco_evaluator.update(res)
evaluator_time = time.time() - evaluator_time
metric_logger.update(model_time=model_time,
evaluator_time=evaluator_time)
# gather the stats from all processes
metric_logger.synchronize_between_processes()
print("Averaged stats:", metric_logger)
coco_evaluator.synchronize_between_processes()
# accumulate predictions from all images
coco_evaluator.accumulate()
coco_evaluator.summarize()
torch.set_num_threads(n_threads)
def get_coloured_mask(mask, pred_cls):
r = np.zeros_like(mask).astype(np.uint8)
g = np.zeros_like(mask).astype(np.uint8)
b = np.zeros_like(mask).astype(np.uint8)
colours = [[0, 0, 0],[0, 255, 0],[0, 255, 255],[255, 255, 0],[80, 70, 180],[180, 40, 250],[245, 145, 50],[70, 150, 250],[50, 190, 190]]
r[mask == 1], g[mask == 1], b[mask == 1] = colours[common.CLASS_NAMES.index(pred_cls)]
coloured_mask = np.stack([r, g, b], axis=2)
return coloured_mask
if __name__ == '__main__':
# load images
img_paths = glob.glob(common.TRAIN_DATA_IMGS)
# load model
model = model.get_model_instance_segmentation(common.NUM_CLASSES)
model, device = common.setup_device(model)
model.load_state_dict(torch.load("model.pth", map_location=device))
model.eval()
confidence = 0.5
for idx in range(len(img_paths)):
# Prediction
img_path = img_paths[idx]
img = Image.open(img_path)
transform = T.Compose([T.ToTensor()])
img = torchvision.transforms.functional.to_tensor(img)
pred = model([img.to(device)])
pred_score = list(pred[0]['scores'].detach().cpu().numpy())
pred_t = [pred_score.index(x) for x in pred_score if x>confidence]
def generate_mask(img_base, img_name):
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
num_classes = 2
image_mean = [
0.34616187074865956, 0.34616187074865956, 0.34616187074865956
]
image_std = [0.10754412766560431, 0.10754412766560431, 0.10754412766560431]
mask_rcnn = get_model_instance_segmentation(num_classes,
image_mean,
image_std,
stats=True)
mask_rcnn.to(device)
mask_rcnn.eval()
model_path = "models/epoch_0099.param"
mask_rcnn.load_state_dict(torch.load(model_path))
# tif files and png files location
img_path = img_base + "/split_pngs/"
tif_path = img_base + "/split_tifs/"
# tmp folder has the vector files created for each craters
# create if it doesn't exist
tmp_folder = img_base + "/split_tifs/tmp/"
if not os.path.exists(tmp_folder):
os.makedirs(tmp_folder)
pred_folder = img_base + "/split_tifs/shape_files/"
if not os.path.exists(pred_folder):
os.makedirs(pred_folder)
print(img_name)
# file name without extension
file_name = img_name.split(".")[0]
individual_folder = tmp_folder + file_name + "/"
if not os.path.exists(individual_folder):
os.makedirs(individual_folder)
image = Image.open(img_path + img_name).convert("RGB")
image = np.array(image)
image = torch.from_numpy(image / 255.0).float()
image = image.permute((2, 0, 1))
pred = mask_rcnn(image.unsqueeze(0).to(device))[0]
#from visualize import display_instances
boxes_ = pred["boxes"].cpu().detach().numpy().astype(int)
boxes = np.empty_like(boxes_)
boxes[:,
0], boxes[:,
1], boxes[:,
2], boxes[:,
3] = boxes_[:,
1], boxes_[:,
0], boxes_[:,
3], boxes_[:,
2]
labels = pred["labels"].cpu().detach().numpy()
scores = pred["scores"].cpu().detach().numpy()
masks = pred["masks"]
indices = scores > 0.6
# Show Scores
boxes = boxes[indices]
labels = labels[indices]
scores = scores[indices]
masks = masks[indices].squeeze(1)
masks = (masks.permute(
(1, 2, 0)).cpu().detach().numpy() > 0.5).astype(np.uint8)
# Show image predictions and save as npy files
image = image.permute((1, 2, 0)).cpu().detach().numpy() * 255
#display_instances(image, boxes, masks, labels, class_names=["background", "crater"], scores=scores)
print(masks.shape)
# Save masks as tif files
crater_prediction_count = masks.shape[2]
for i in range(crater_prediction_count):
# create single crater tif file
tmp_tif = individual_folder + file_name + "_" + str(i) + ".tif"
os.system("cp " + tif_path + file_name + ".tif " + tmp_tif)
dataset = gdal.Open(tmp_tif, gdal.GA_Update)
band = dataset.GetRasterBand(1)
new_mask = masks[:, :, i]
new_mask = new_mask
band.SetNoDataValue(0)
band.WriteArray(new_mask)
dataset.FlushCache()
del dataset
# convert the tif file to a vector data gdal polygonize
tmp_shp = individual_folder + file_name + "_" + str(i) + ".shp"
os.system(
"""gdal_polygonize.py {} {} -b 1 -f "ESRI Shapefile" 1 DN """.
format(tmp_tif, tmp_shp))
if crater_prediction_count > 0:
os.system("ogrmerge.py -single -overwrite_ds -o {} {}*.shp".format(
pred_folder + file_name + ".shp", individual_folder))
os.system("rm -rf " + individual_folder + "*")
def test(config):
# test_dt = FashionDataset(config, transforms= None)
sample_df = pd.read_csv(config.sample_path)
################################################################################
# create the model instance
num_classes = 46 + 1
model_test = get_model_instance_segmentation(num_classes)
#load the training weights
# load_path =osp.join(config.save_dir, '9_weights')
load_path = osp.join(config.save_dir, 'weights')
model_test.load_state_dict(torch.load(osp.join(load_path, '9_model.bin')))
# send the test model to gpu
model_test.to(device)
for param in model_test.parameters():
param.requires_grad = False
model_test.eval()
# for submission
sub_list = []
missing_count = 0
for i, row in tqdm(sample_df.iterrows(), total=len(sample_df)):
###modify##########################################################
# import the image
img_path = osp.join(config.test_dir, sample_df['ImageId'][i])
img = Image.open(img_path).convert('RGB')
img = img.resize((config.width, config.height),
resample=Image.BILINEAR)
# convert the img as tensor
img = F.to_tensor(img)
#####modify#############################################################
pred = model_test([img.to(device)])[0]
masks = np.zeros((512, 512, len(pred['masks'])))
for j, m in enumerate(pred['masks']):
res = transforms.ToPILImage()(m.permute(1, 2, 0).cpu().numpy())
res = np.asarray(res.resize((512, 512), resample=Image.BILINEAR))
masks[:, :, j] = (res[:, :] * 255. > 127).astype(np.uint8)
labels = pred['labels'].cpu().numpy()
scores = pred['scores'].cpu().numpy()
#
best_idx = 0
# print('the maximum scores is {}'.format(np.mean(scores)))
# print('the current masks is {}'.format(masks))
for _scores in scores:
if _scores > 0.8:
best_idx += 1
if best_idx == 0:
# print(masks.shape[-1])
sub_list.append([sample_df.loc[i, 'ImageId'], '1 1', 23])
missing_count += 1
continue
if masks.shape[-1] > 0:
masks = refine_masks(masks[:, :, :best_idx], labels[:best_idx])
for m, rle, label in masks:
sub_list.append([sample_df.loc[i, 'ImageId'], rle, label])
else:
sub_list.append([sample_df.loc[i, 'ImageId'], '1 1', 23])
missing_count += 1
submission_df = pd.DataFrame(sub_list, columns=sample_df.columns.values)
print("Total image results: ", submission_df['ImageId'].nunique())
print("Missing Images: ", missing_count)
submission_df = submission_df[submission_df.EncodedPixels.notnull()]
for row in range(len(submission_df)):
line = submission_df.iloc[row, :]
submission_df.iloc[row, 1] = line['EncodedPixels'].replace('.0', '')
# submission_df.head()
submit_path = config.submit_path + r'submission.csv'
print('submit_path: ', submit_path)
submission_df.to_csv(submit_path, index=False)
print('ok,finished')
def main():
parser = argparse.ArgumentParser(description='Mask R-CNN')
parser.add_argument('--dataset',
default='unimib',
type=str,
choices=['unimib', 'food201', 'chfood'],
help='dataset name')
args = parser.parse_args()
# train on the GPU or on the CPU, if a GPU is not available
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
if args.dataset == 'unimib':
num_classes = 74
# use our dataset and defined transformations
dataset = UNIMIBDataset(
'/home/hatsunemiku/dev/mask-rcnn/data/UNIMIB2016',
get_transform(train=True))
dataset_test = UNIMIBDataset(
'/home/hatsunemiku/dev/mask-rcnn/data/UNIMIB2016',
get_transform(train=False), False)
elif args.dataset == 'food201':
num_classes = 2
# use our dataset and defined transformations
dataset = Food201Dataset(
'/home/hatsunemiku/dev/mask-rcnn/data/food201',
get_transform(train=True))
dataset_test = Food201Dataset(
'/home/hatsunemiku/dev/mask-rcnn/data/food201',
get_transform(train=False), False)
elif args.dataset == 'chfood':
num_classes = 24
# use our dataset and defined transformations
dataset = CHFoodDataset('/home/hatsunemiku/dev/mask-rcnn/data/ch_food',
get_transform(train=True))
dataset_test = CHFoodDataset(
'/home/hatsunemiku/dev/mask-rcnn/data/ch_food',
get_transform(train=False), False)
else:
raise Exception()
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=2,
shuffle=True,
num_workers=0,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(dataset_test,
batch_size=1,
shuffle=False,
num_workers=0,
collate_fn=utils.collate_fn)
# get the model using our helper function
model = get_model_instance_segmentation(num_classes)
# move model to the right device
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.002,
momentum=0.9,
weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=10,
gamma=0.1)
# let's train it for 10 epochs
num_epochs = 30
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model,
optimizer,
data_loader,
device,
epoch,
print_freq=10)
print('epoch %d, lr -----------> %.6f' %
(epoch, optimizer.param_groups[0]['lr']))
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, data_loader_test, device=device)
torch.save(model.state_dict(), 'chfood_model.pkl')
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices)
indices_test = torch.randperm(len(dataset_test)).tolist()
dataset_test = torch.utils.data.Subset(dataset_test, indices_test)
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=4, shuffle=True, num_workers=4,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1, shuffle=False, num_workers=4,
collate_fn=utils.collate_fn)
# get the model using our helper function
mask_rcnn = get_model_instance_segmentation(num_classes, None, None, False)
read_param = False
if read_param:
mask_rcnn.load_state_dict(torch.load("trained_param_eureka_aug_bin/epoch_0009.param"))
# move model to the right device
mask_rcnn.to(device)
# construct an optimizer
params = [p for p in mask_rcnn.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.01,
momentum=0.9, weight_decay=0.00001)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=10,
def main():
# mode argument
args = argparse.ArgumentParser()
args.add_argument("--num_classes", type=int, default=22)
args.add_argument("--lr", type=int, default=0.005)
args.add_argument("--cuda", type=bool, default=True)
args.add_argument("--num_epochs", type=int, default=30)
args.add_argument("--model_name", type=str, default="mask_1.pth")
args.add_argument("--prediction_file", type=str, default="prediction")
args.add_argument("--batch", type=int, default=16)
args.add_argument("--mode", type=str, default="train")
config = args.parse_args()
num_classes = config.num_classes
base_lr = config.lr
cuda = config.cuda
num_epochs = config.num_epochs
model_name = config.model_name
prediction_file = config.prediction_file
batch = config.batch
mode = config.mode
# 도움 함수를 이용해 모델을 가져옵니다
new_model = model.get_model_instance_segmentation(num_classes)
# 학습을 GPU로 진행하되 GPU가 가용하지 않으면 CPU로 합니다
device = torch.device('cuda') if cuda else torch.device('cpu')
# 모델을 GPU나 CPU로 옮깁니다
new_model.to(device)
if mode == 'train':
# 데이터셋과 정의된 변환들을 사용합니다
'''dataset = CustomDataset(DATASET_PATH, dataloader.get_transform(train=True))
dataset_val = CustomDataset(DATASET_PATH, dataloader.get_transform(train=False))
# 데이터셋을 학습용과 테스트용으로 나눕니다(역자주: 여기서는 전체의 50개를 테스트에, 나머지를 학습에 사용합니다)
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-10])
dataset_val = torch.utils.data.Subset(dataset_val, indices[-10:])
data_loader_val = torch.utils.data.DataLoader(
dataset_val, batch_size=batch, shuffle=True,num_workers=4, collate_fn=dataloader.collate_fn)
'''
dataset =dataloader.make_dataset(DATASET_PATH)
# 데이터 로더를 학습용과 검증용으로 정의합니다
dataset_loader = torch.utils.data.DataLoader(
dataset, batch_size=batch, shuffle=True, num_workers=4,collate_fn=dataloader.collate_fn)
# 옵티마이저(Optimizer)를 만듭니다
params = [p for p in new_model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=base_lr,
momentum=0.9, weight_decay=0.0005)
# 학습률 스케쥴러를 만듭니다
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.1)
train(new_model, dataset_loader, device, num_epochs, optimizer=optimizer, lr_scheduler=lr_scheduler)
elif mode == 'test' :
dataset_test =dataloader.make_dataset(DATASET_PATH)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1, shuffle=False, collate_fn=dataloader.collate_fn)
load_model(model_name, new_model)
test(new_model, data_loader_test, device, prediction_file)
print("That's it!")
trainset = torch.utils.data.Subset(trainset, indeces[:-50])
testset = torch.utils.data.Subset(testset, indeces[-50:])
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=2,
num_workers=4,
shuffle=True,
collate_fn=utils.collate_fn)
testloader = torch.utils.data.DataLoader(testset,
batch_size=1,
num_workers=4,
shuffle=True,
collate_fn=utils.collate_fn)
num_classes = 2
model = get_model_instance_segmentation(num_classes)
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=5e-3,
momentum=0.9, weight_decay=5e-4)
lr_schedular = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
epochs = 10
for epoch in range(epochs):
train_one_epoch(model, optimizer, trainloader, device, epoch, print_freq=10)
lr_schedular.step()
evaluate(model, testloader, device)
def __init__(self):
super(Ui_MainWindow, self).__init__()
self.num_classes = 46 + 1
self.model_test = get_model_instance_segmentation(self.num_classes)
self.model_path = './9_model.bin'
self.pretrain_params = torch.load(self.model_path, map_location='cpu')
for k in list(self.pretrain_params.keys()):
if k.startswith('module.'):
self.pretrain_params[
k[len('module.'):]] = self.pretrain_params[k]
del self.pretrain_params[k]
self.model_test.load_state_dict(self.pretrain_params)
# send the test model to gpu
self.model_test.to(device)
for param in self.model_test.parameters():
param.requires_grad = False
self.model_test.eval()
#self.config, self.unparsed = get_config()
self.image_name = None
self.raw_image = None
self.logo_size = (350, 350)
self.mask_size = (60, 60)
self.width = 512
self.height = 512
self.pred = None
self.classes = None
self.save_path = './images_pred/'
with open('label_descriptions.json') as json_f:
look_up_classes = json.load(json_f)
self.categories = look_up_classes["categories"]
self.attributes = look_up_classes["attributes"]
if not os.path.exists('./images_pred/'):
os.makedirs('./images_pred/')
self.dic = {
#'003b3e4579833df109d330773418fd65':{33:[182],31:[160,204],24:[],10:[108,141,155]},
'008324a52c5910d338807fb785ebaa0a': {
13: [],
10: [101, 108, 128, 142],
31: [158, 204],
28: [163],
23: []
},
'wen': {
33: [183],
9: [91, 115, 136, 145],
29: [174],
32: [218],
31: [160]
},
'3people': {
1: [0, 115, 145],
31: [157, 204],
24: [],
33: [182]
},
#'00beb7542b35409dd316572e497e701c':{8:[65,115,118,229],1:[11,142],33:[],24:[]}
}
self.setupUi()
