# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import six
import sys
sys.path.append('../../')

from models.resnet import rf_lw50

from utils.helpers import prepare_img

import glob

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from skimage import io, color, img_as_uint, img_as_ubyte
from PIL import Image
import pydensecrf.densecrf as dcrf
from scipy.special import softmax

from pydensecrf.utils import compute_unary, create_pairwise_bilateral, \
    create_pairwise_gaussian, unary_from_softmax
# -

has_cuda = torch.cuda.is_available()
n_classes = 3
thresh = 0.8

net = nn.DataParallel(rf_lw50(num_classes=3, pretrained=True))
weights_path = "checkpoints/checkpoint_9218.pth.tar"
net.load_state_dict(torch.load(weights_path)["model"])

# read image
frame = Image.open("/home/rafay/Pictures/Webcam/2022-01-19-225013.jpg")

frame





# convert frame to np array
frame = np.array(frame)
print(frame.dtype)
print(frame.shape)

img_inp = torch.tensor(prepare_img(frame).transpose(2, 0, 1)[None]).float()

# +
segm = net(img_inp)[0].data.cpu().numpy().transpose(1, 2, 0)
segm = softmax(segm, axis=2)
softmax_out = np.copy(segm)  # copy for later
segm = cv2.resize(segm, (640,480), interpolation=cv2.INTER_CUBIC)
print(segm.shape) # 120,160,3
for channel in range(3):
    grid = segm[:,:,channel]
    grid[grid >= thresh] = 1.0
    grid[grid < thresh] = 0.0
    segm[:, :, channel] = grid

# segm = cmap[segm.argmax(axis=2).astype(np.uint8)]
segm = segm.argmax(axis=2).astype(np.uint8)
final_img = np.array(color.label2rgb(np.array(segm), frame, colors=[(255, 0, 0), (0, 0, 255)],alpha=0.01, bg_label=0, bg_color=None))
plt.imshow(final_img)

# +
softmax_out = cv2.resize(softmax_out, (640,480), interpolation=cv2.INTER_CUBIC)

feat_first = softmax_out.transpose((2, 0, 1)).reshape((3,-1))  # reshape to n_classes, -1
print(feat_first.shape)
unary = unary_from_softmax(feat_first)
unary = np.ascontiguousarray(unary)
# -

d = dcrf.DenseCRF2D(frame.shape[1], frame.shape[0], n_classes)  # 640x480, 3 classes

d.setUnaryEnergy(unary)
d.addPairwiseGaussian(sxy=(5, 5), compat=3, kernel=dcrf.DIAG_KERNEL,
                          normalization=dcrf.NORMALIZE_SYMMETRIC)

d.addPairwiseBilateral(sxy=(5, 5), srgb=(13, 13, 13), rgbim=frame,
                       compat=10,
                       kernel=dcrf.DIAG_KERNEL,
                       normalization=dcrf.NORMALIZE_SYMMETRIC)

Q = d.inference(5)

res = np.argmax(Q, axis=0).reshape((frame.shape[0], frame.shape[1]))

print(res.shape)

print(np.unique(res))



final_img = np.array(color.label2rgb(np.array(res), frame, colors=[(255, 0, 0), (0, 0, 255)],alpha=0.01, bg_label=0, bg_color=None))

plt.figure(figsize=(10,10))
plt.imshow(final_img)
plt.show()



StarcoderdataJupyterScriptsDedupFiltered