Spaces:
Sleeping
Sleeping
File size: 3,753 Bytes
1ed7deb |
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 |
import sys, os
import numpy as np
import scipy
import torch
import torch.nn as nn
from scipy import ndimage
from tqdm import tqdm, trange
from PIL import Image
import torch.hub
import torchvision
import torch.nn.functional as F
# download deeplabv2_resnet101_msc-cocostuff164k-100000.pth from
# https://github.com/kazuto1011/deeplab-pytorch/releases/download/v1.0/deeplabv2_resnet101_msc-cocostuff164k-100000.pth
# and put the path here
CKPT_PATH = "TODO"
rescale = lambda x: (x + 1.) / 2.
def rescale_bgr(x):
x = (x+1)*127.5
x = torch.flip(x, dims=[0])
return x
class COCOStuffSegmenter(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.n_labels = 182
model = torch.hub.load("kazuto1011/deeplab-pytorch", "deeplabv2_resnet101", n_classes=self.n_labels)
ckpt_path = CKPT_PATH
model.load_state_dict(torch.load(ckpt_path))
self.model = model
normalize = torchvision.transforms.Normalize(mean=self.mean, std=self.std)
self.image_transform = torchvision.transforms.Compose([
torchvision.transforms.Lambda(lambda image: torch.stack(
[normalize(rescale_bgr(x)) for x in image]))
])
def forward(self, x, upsample=None):
x = self._pre_process(x)
x = self.model(x)
if upsample is not None:
x = torch.nn.functional.upsample_bilinear(x, size=upsample)
return x
def _pre_process(self, x):
x = self.image_transform(x)
return x
@property
def mean(self):
# bgr
return [104.008, 116.669, 122.675]
@property
def std(self):
return [1.0, 1.0, 1.0]
@property
def input_size(self):
return [3, 224, 224]
def run_model(img, model):
model = model.eval()
with torch.no_grad():
segmentation = model(img, upsample=(img.shape[2], img.shape[3]))
segmentation = torch.argmax(segmentation, dim=1, keepdim=True)
return segmentation.detach().cpu()
def get_input(batch, k):
x = batch[k]
if len(x.shape) == 3:
x = x[..., None]
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
return x.float()
def save_segmentation(segmentation, path):
# --> class label to uint8, save as png
os.makedirs(os.path.dirname(path), exist_ok=True)
assert len(segmentation.shape)==4
assert segmentation.shape[0]==1
for seg in segmentation:
seg = seg.permute(1,2,0).numpy().squeeze().astype(np.uint8)
seg = Image.fromarray(seg)
seg.save(path)
def iterate_dataset(dataloader, destpath, model):
os.makedirs(destpath, exist_ok=True)
num_processed = 0
for i, batch in tqdm(enumerate(dataloader), desc="Data"):
try:
img = get_input(batch, "image")
img = img.cuda()
seg = run_model(img, model)
path = batch["relative_file_path_"][0]
path = os.path.splitext(path)[0]
path = os.path.join(destpath, path + ".png")
save_segmentation(seg, path)
num_processed += 1
except Exception as e:
print(e)
print("but anyhow..")
print("Processed {} files. Bye.".format(num_processed))
from taming.data.sflckr import Examples
from torch.utils.data import DataLoader
if __name__ == "__main__":
dest = sys.argv[1]
batchsize = 1
print("Running with batch-size {}, saving to {}...".format(batchsize, dest))
model = COCOStuffSegmenter({}).cuda()
print("Instantiated model.")
dataset = Examples()
dloader = DataLoader(dataset, batch_size=batchsize)
iterate_dataset(dataloader=dloader, destpath=dest, model=model)
print("done.")
|