mirror of
https://github.com/wassname/PSPNet-Keras-tensorflow.git
synced 2026-08-21 11:11:01 +08:00
76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
from PIL import Image, ImageDraw
|
|
import scipy.ndimage
|
|
import scipy.io
|
|
import numpy as np
|
|
import time
|
|
import copy
|
|
|
|
|
|
class BaseDraw:
|
|
def __init__(self, color150, objectNames, img, pred_size, predicted_classes):
|
|
self.class_colors = scipy.io.loadmat(color150)
|
|
self.class_names = scipy.io.loadmat(objectNames, struct_as_record=False)
|
|
self.im = img
|
|
self.pred_size = pred_size
|
|
self.predicted_classes = copy.deepcopy(predicted_classes)
|
|
self.original_W = self.im.size[0]
|
|
self.original_H = self.im.size[1]
|
|
|
|
self.output_W = self.original_W
|
|
self.output_H = self.original_H
|
|
|
|
|
|
def dumpArray(self, array, i):
|
|
test = array*100
|
|
test = Image.fromarray(test.astype('uint8'))
|
|
test = test.convert("RGB")
|
|
test.save('/home/vlad/oS_AI/'+str(i)+'t.jpg', "JPEG")
|
|
|
|
def calculateResize(self):
|
|
W_coef = float(self.original_W)/float(self.output_W)
|
|
H_coef = float(self.original_H)/float(self.output_H)
|
|
horiz_pad = 0
|
|
vert_pad = 0
|
|
if W_coef > H_coef:
|
|
coef = W_coef
|
|
horiz_pad = int((self.output_H - self.original_H/coef)/2)
|
|
return [coef, horiz_pad, vert_pad]
|
|
else:
|
|
coef = H_coef
|
|
vert_pad = int((self.output_W - self.original_W/coef)/2)
|
|
return [coef, horiz_pad, vert_pad]
|
|
|
|
|
|
def resizeToOutput(self, image, coef, h_pad, w_pad):
|
|
image = image.resize((int(self.original_W/coef), int(self.original_H/coef)), resample=Image.BILINEAR)
|
|
outputImage = Image.new("RGB",(self.output_W,self.output_H),(0,0,0))
|
|
outputImage.paste(image,(w_pad,h_pad))
|
|
return outputImage
|
|
|
|
|
|
|
|
def drawSimpleSegment(self):
|
|
|
|
#Drawing module
|
|
im_Width, im_Height = self.pred_size
|
|
prediction_image = Image.new("RGB", (im_Width, im_Height) ,(0,0,0))
|
|
prediction_imageDraw = ImageDraw.Draw(prediction_image)
|
|
|
|
#BASE all image segmentation
|
|
for i in range(im_Width):
|
|
for j in range(im_Height):
|
|
#get matrix element class(0-149)
|
|
px_Class = self.predicted_classes[j][i]
|
|
#assign color from .mat list
|
|
put_Px_Color = tuple(self.class_colors['colors'][px_Class])
|
|
|
|
#drawing
|
|
prediction_imageDraw.point((i,j), fill=put_Px_Color)
|
|
|
|
#Resize to original size and save
|
|
self.coef, self.h_pad, self.w_pad = self.calculateResize()
|
|
FullHdOutImage = self.resizeToOutput(prediction_image, self.coef, self.h_pad, self.w_pad)
|
|
FullHdOutImage = Image.blend(FullHdOutImage, self.im, 0.5)
|
|
|
|
return FullHdOutImage
|