-
Notifications
You must be signed in to change notification settings - Fork 30
/
Copy pathvisualizeImg.py
100 lines (87 loc) · 3.46 KB
/
visualizeImg.py
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
from PIL import Image, ImageDraw, ImageFont, ImageFilter
import matplotlib.pyplot as plt
import matplotlib
import numpy
import random
import torch
import os
from utils import *
from torchvision import transforms as transforms
RECTANGLE_SIZE = 4
imgId = 0
# random color
def randomColor():
return (random.randint(64, 255), random.randint(64, 255), random.randint(64, 255))
# input: numpy.ndarray
# topleft_Corner: (x,y)
# width: rectangle width
# height: rectangle height
# size: line size
def drawRectangle(input, x, y, width, height, size, k, className, confidence):
# numpy->PIL
inputImg = Image.fromarray(numpy.uint8(input))
draw = ImageDraw.Draw(inputImg)
color = randomColor()
font = ImageFont.truetype('Arial.ttf', 24)
for i in range(1, size + 1):
draw.rectangle((x + (size - i), y + (size - i), x + width + i, y + height + i), outline=color)
draw.text((x+size+6,y), k, font = font, fill = color)
draw.text((x+size+30,y), className, font = font, fill = color)
draw.text((x+len(className)*14+36,y), confidence, font = font, fill = color)
return inputImg
# M: k*batchsize*2*3
# sourceCoordinate: k*batchsize*2*2
# return: topleft, bottomright corners
def getSourceCoordinate(M):
target = torch.tensor([-1.,1.,1.,-1.,1.,1.])
target = target.view(3,2)
sourceCoordinate = torch.matmul(M,target)
if imgId == 0:
for batch_index in range(M.size(1)):
print("img", batch_index,"-----------")
for k_index in range(M.size(0)):
print("k: ", k_index)
print(M[k_index, batch_index, :, :])
x0 = sourceCoordinate[:,:,0,0]*256.+256.
x1 = sourceCoordinate[:,:,0,1]*256.+256.
y0 = -(sourceCoordinate[:,:,1,0]*256.-256.)
y1 = -(sourceCoordinate[:,:,1,1]*256.-256.)
sourceCoordinate[:,:,0,0] = torch.min(x0, x1)
sourceCoordinate[:,:,0,1] = torch.max(x0, x1)
sourceCoordinate[:,:,1,0] = torch.min(y0, y1)
sourceCoordinate[:,:,1,1] = torch.max(y0, y1)
return sourceCoordinate
# sourceCoordinate: k*batchsize*2*2
# rectangleInfo: k*batchsize*4
# return:k*batchsize*[x,y,width,height]
def getPredictedRectangle(M):
coordinate = getSourceCoordinate(M)
rectangleInfo = torch.zeros(coordinate.size(0), coordinate.size(1), 4)
rectangleInfo[:,:,0] = coordinate[:,:,0,0]
rectangleInfo[:,:,1] = coordinate[:,:,1,0]
rectangleInfo[:,:,2] = coordinate[:,:,0,1]-coordinate[:,:,0,0]
rectangleInfo[:,:,3] = coordinate[:,:,1,1]-coordinate[:,:,1,0]
return rectangleInfo
def drawPictures(original_img, rectangle, className, confidence):
for batch_index in range(shape(className)[1]):
current_picture = transforms.ToPILImage()(original_img[batch_index]).convert('RGB')
for iterator_index in range(shape(className)[0]):
k = str(iterator_index)
x = rectangle[iterator_index, batch_index, 0]
y = rectangle[iterator_index, batch_index, 1]
width = rectangle[iterator_index, batch_index, 2]
height = rectangle[iterator_index, batch_index, 3]
classname = className[iterator_index, batch_index]
conf_value = str(round(confidence[iterator_index, batch_index].item(),3))
current_picture = drawRectangle(current_picture, x, y, width, height, RECTANGLE_SIZE, k, classname, conf_value)
global imgId
current_picture.save("./result_visual/visualize_imgs_with_priori/"+str(imgId)+".jpg")
imgId += 1
return
def visualize_attentional_regions(original_img, M, scores):
# rectangle
rectangle = getPredictedRectangle(M)
# information
confidence, className = getPredictedInfo(scores)
drawPictures(original_img, rectangle, className, confidence)
return