<jupyter_start><jupyter_text>NYUD-three: semantic segmentation, depth estimation, surface normals estimation<jupyter_code>%matplotlib inline

import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import sys
sys.path.append('../')
from models import net
import cv2
import torch
from torch.autograd import Variable
# Pre-processing and post-processing constants #
CMAP = np.load('../cmap_nyud.npy')
DEPTH_COEFF = 5000. # to convert into metres
HAS_CUDA = torch.cuda.is_available()
IMG_SCALE  = 1./255
IMG_MEAN = np.array([0.485, 0.456, 0.406]).reshape((1, 1, 3))
IMG_STD = np.array([0.229, 0.224, 0.225]).reshape((1, 1, 3))
MAX_DEPTH = 8.
MIN_DEPTH = 0.
NUM_CLASSES = 40
NUM_TASKS = 3 # segm + depth + normals
def prepare_img(img):
    return (img * IMG_SCALE - IMG_MEAN) / IMG_STD
model = net(num_classes=NUM_CLASSES, num_tasks=NUM_TASKS)
if HAS_CUDA:
    _ = model.cuda()
_ = model.eval()
ckpt = torch.load('../../weights/ExpNYUD_three.ckpt')
model.load_state_dict(ckpt['state_dict'])
# Figure 5-bottom row
img_path = '../../examples/ExpNYUD_three/000433.png'
img = np.array(Image.open(img_path))
gt_segm = np.array(Image.open('../../examples/ExpNYUD_three/segm_gt_000433.png'))
gt_depth = np.array(Image.open('../../examples/ExpNYUD_three/depth_gt_000433.png'))
gt_norm = np.array(Image.open('../../examples/ExpNYUD_three/norm_gt_000433.png'))
with torch.no_grad():
    img_var = Variable(torch.from_numpy(prepare_img(img).transpose(2, 0, 1)[None]), requires_grad=False).float()
    if HAS_CUDA:
        img_var = img_var.cuda()
    segm, depth, norm = model(img_var)
    segm = cv2.resize(segm[0, :NUM_CLASSES].cpu().data.numpy().transpose(1, 2, 0),
                      img.shape[:2][::-1],
                      interpolation=cv2.INTER_CUBIC)
    depth = cv2.resize(depth[0, 0].cpu().data.numpy(),
                       img.shape[:2][::-1],
                       interpolation=cv2.INTER_CUBIC)
    norm = cv2.resize(norm[0].cpu().data.numpy().transpose(1, 2, 0),
                       img.shape[:2][::-1],
                       interpolation=cv2.INTER_CUBIC)
    segm = CMAP[segm.argmax(axis=2) + 1].astype(np.uint8)
    depth = np.abs(depth)
    out_norm = norm / np.linalg.norm(norm, axis=2, keepdims=True)
    ## xzy->RGB ##
    out_norm[:, :, 0] = ((out_norm[:, :, 0] + 1.) / 2.) * 255.
    out_norm[:, :, 1] = ((out_norm[:, :, 1] + 1.) / 2.) * 255.
    out_norm[:, :, 2] = ((1. - out_norm[:, :, 2]) / 2.) * 255.
    out_norm = out_norm.astype(np.uint8)

plt.figure(figsize=(18, 12))
plt.subplot(171)
plt.imshow(img)
plt.title('orig img')
plt.axis('off')
plt.subplot(172)
plt.imshow(CMAP[gt_segm + 1])
plt.title('gt segm')
plt.axis('off')
plt.subplot(173)
plt.imshow(segm)
plt.title('pred segm')
plt.axis('off')
plt.subplot(174)
plt.imshow(gt_depth / DEPTH_COEFF, cmap='plasma', vmin=MIN_DEPTH, vmax=MAX_DEPTH)
plt.title('gt depth')
plt.axis('off')
plt.subplot(175)
plt.imshow(depth, cmap='plasma', vmin=MIN_DEPTH, vmax=MAX_DEPTH)
plt.title('pred depth')
plt.axis('off')
plt.subplot(176)
plt.imshow(gt_norm)
plt.title('gt norm')
plt.axis('off')
plt.subplot(177)
plt.imshow(out_norm)
plt.title('pred norm')
plt.axis('off');<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18404522613065327, 'fraction_numerical': 0.04585427135678392, 'mean_word_length': 3.789473684210526, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '426207', 'license': 'BSD-2-Clause', 'n_tokens_mistral': 1413, 'n_tokens_neox': 1273, 'n_words': 217, 'path': 'src/notebooks/ExpNYUD_three.ipynb', 'repo_name': 'chihong0522/multi-task-refinenet', 'size': '3184'}