File size: 4,670 Bytes
7754b29 |
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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import json
from torch.utils import data
from torchvision.datasets import ImageFolder
import torch
import os
from PIL import Image
import numpy as np
import argparse
from tqdm import tqdm
from munkres import Munkres
import multiprocessing
from multiprocessing import Process, Manager
import collections
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
import random
import torchvision
import cv2
import random
torch.manual_seed(0)
SegItem = collections.namedtuple('SegItem', ('image_name', 'tag'))
normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
TRANSFORM_TRAIN = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
])
TRANSFORM_EVAL = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
])
IMAGE_TRANSFORMS = transforms.Compose([
transforms.ToTensor(),
normalize
])
MERGED_TAGS = {'n04356056', 'n04355933',
'n04493381', 'n02808440',
'n03642806', 'n03832673',
'n04008634', 'n03773504',
'n03887697', 'n15075141'}
TRAIN_PARTITION = "train"
VAL_PARTITION = "val"
LEGAL_PARTITIONS = {TRAIN_PARTITION, VAL_PARTITION}
# TRAIN_CLASSES = 500
class SegmentationDataset(ImageFolder):
def __init__(self, seg_path, imagenet_path, partition=TRAIN_PARTITION, num_samples=2, train_classes=500
, imagenet_classes_path='imagenet_classes.json', seed=None):
assert partition in LEGAL_PARTITIONS
self._partition = partition
self._seg_path = seg_path
self._imagenet_path = imagenet_path
with open(imagenet_classes_path, 'r') as f:
self._imagenet_classes = json.load(f)
self._tag_list = [tag for tag in os.listdir(self._seg_path) if tag not in MERGED_TAGS]
if seed:
print(f'Shuffling training classes with seed {seed}')
random.seed(seed)
random.shuffle(self._tag_list)
if partition == TRAIN_PARTITION:
# Skip merged tags
self._tag_list = self._tag_list[:train_classes]
elif partition == VAL_PARTITION:
# Skip merged tags
self._tag_list = self._tag_list[train_classes:]
for tag in self._tag_list:
assert tag in self._imagenet_classes
self._all_segementations = []
for tag in self._tag_list:
base_dir = os.path.join(self._seg_path, tag)
for i, seg in enumerate(os.listdir(base_dir)):
if i >= num_samples:
break
self._all_segementations.append(SegItem(seg.split('.')[0], tag))
def __getitem__(self, item):
seg_item = self._all_segementations[item]
seg_path = os.path.join(self._seg_path, seg_item.tag, seg_item.image_name + ".png")
image_path = os.path.join(self._imagenet_path, seg_item.tag, seg_item.image_name + ".JPEG")
seg_map = Image.open(seg_path)
image = Image.open(image_path)
image = image.convert('RGB')
seg_map = np.array(seg_map)
seg_map = seg_map[:, :, 1] * 256 + seg_map[:, :, 0]
assert len([cand for cand in np.unique(seg_map) if cand != 0 and cand != 1000]) == 1
# Convert to binary seg maps
seg_map[seg_map == 1000] = 0
seg_map[seg_map != 0] = 1
seg_map = torch.from_numpy(seg_map.astype(np.float32))
# transforms - start
seg_map = seg_map.reshape(1, seg_map.shape[-2], seg_map.shape[-1])
if self._partition == VAL_PARTITION:
image = TRANSFORM_EVAL(image)
seg_map = TRANSFORM_EVAL(seg_map)
elif self._partition == TRAIN_PARTITION:
# Resize
resize = transforms.Resize(size=(256, 256))
image = resize(image)
seg_map = resize(seg_map)
# Random crop
i, j, h, w = transforms.RandomCrop.get_params(
image, output_size=(224, 224))
image = TF.crop(image, i, j, h, w)
seg_map = TF.crop(seg_map, i, j, h, w)
# RandomHorizontalFlip
if random.random() > 0.5:
image = TF.hflip(image)
seg_map = TF.hflip(seg_map)
else:
raise Exception(f"Unsupported partition type {self._partition}")
# normalize original image and turn to tensor
image_ten = IMAGE_TRANSFORMS(image)
# transforms - end
class_name = int(self._imagenet_classes[seg_item.tag])
return seg_map, image_ten, class_name
def __len__(self):
return len(self._all_segementations) |