update utils

utils/convlstm.py

@@ -2,6 +2,7 @@
 from torch import nn
 from torch.autograd import Variable
 
+
 class ConvLSTMCell(nn.Module):
 
     def __init__(self, input_size, input_dim, hidden_dim, kernel_size, bias):
@@ -41,10 +42,12 @@ def __init__(self, input_size, input_dim, hidden_dim, kernel_size, bias):
     def forward(self, input_tensor, cur_state):
         h_cur, c_cur = cur_state
 
-        combined = torch.cat([input_tensor, h_cur], dim=1)  # concatenate along channel axis
+        combined = torch.cat([input_tensor, h_cur],
+                             dim=1)  # concatenate along channel axis
 
         combined_conv = self.conv(combined)
-        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
+        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim,
+                                             dim=1)
         i = torch.sigmoid(cc_i)
         f = torch.sigmoid(cc_f)
         o = torch.sigmoid(cc_o)
@@ -56,19 +59,23 @@ def forward(self, input_tensor, cur_state):
         return h_next, c_next
 
     def init_hidden(self, batch_size):
-        return (Variable(torch.zeros(batch_size, self.hidden_dim, self.height, self.width)).cuda(),
-                Variable(torch.zeros(batch_size, self.hidden_dim, self.height, self.width)).cuda())
+        return (Variable(torch.zeros(batch_size, self.hidden_dim, self.height,
+                                     self.width)).cuda(),
+                Variable(torch.zeros(batch_size, self.hidden_dim, self.height,
+                                     self.width)).cuda())
 
 
 class ConvLSTM(nn.Module):
 
-    def __init__(self, input_size, input_dim, hidden_dim, kernel_size, num_layers,
-                 batch_first = False, bias = True, return_all_layers = False):
+    def __init__(self, input_size, input_dim, hidden_dim, kernel_size,
+                 num_layers,
+                 batch_first=False, bias=True, return_all_layers=False):
         super(ConvLSTM, self).__init__()
 
         self._check_kernel_size_consistency(kernel_size)
 
-        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers
+        # Make sure that both `kernel_size` and `hidden_dim` are lists
+        # having len == num_layers
         kernel_size = self._extend_for_multilayer(kernel_size, num_layers)
         hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)
         if not len(kernel_size) == len(hidden_dim) == num_layers:
@@ -86,7 +93,8 @@ def __init__(self, input_size, input_dim, hidden_dim, kernel_size, num_layers,
 
         cell_list = []
         for i in range(0, self.num_layers):
-            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]
+            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[
+                i - 1]
 
             cell_list.append(ConvLSTMCell(input_size=(self.height, self.width),
                                           input_dim=cur_input_dim,
@@ -96,8 +104,7 @@ def __init__(self, input_size, input_dim, hidden_dim, kernel_size, num_layers,
 
         self.cell_list = nn.ModuleList(cell_list)
 
-
-    def forward(self, input_tensor, hidden_state = None):
+    def forward(self, input_tensor, hidden_state=None):
         """
 
         Parameters
@@ -134,8 +141,9 @@ def forward(self, input_tensor, hidden_state = None):
             for t in range(seq_len):
                 #                 print(cur_layer_input.shape)
 
-                h, c = self.cell_list[layer_idx](input_tensor=cur_layer_input[:, t, :, :, :],
-                                                 cur_state=[h, c])
+                h, c = self.cell_list[layer_idx](
+                    input_tensor=cur_layer_input[:, t, :, :, :],
+                    cur_state=[h, c])
                 output_inner.append(h)
 
             layer_output = torch.stack(output_inner, dim=1)
@@ -150,23 +158,21 @@ def forward(self, input_tensor, hidden_state = None):
 
         return layer_output_list, last_state_list
 
-
     def _init_hidden(self, batch_size):
         init_states = []
         for i in range(self.num_layers):
             init_states.append(self.cell_list[i].init_hidden(batch_size))
         return init_states
 
-
     @staticmethod
     def _check_kernel_size_consistency(kernel_size):
         if not (isinstance(kernel_size, tuple) or
-                (isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):
+                (isinstance(kernel_size, list) and all(
+                    [isinstance(elem, tuple) for elem in kernel_size]))):
             raise ValueError('`kernel_size` must be tuple or list of tuples')
 
-
     @staticmethod
     def _extend_for_multilayer(param, num_layers):
         if not isinstance(param, list):
             param = [param] * num_layers
-        return param
+        return param

utils/utils.py

@@ -0,0 +1,76 @@
+from PIL import Image, ImageDraw
+import numpy as np
+import torch
+
+def iou(vertices1, vertices2, h ,w):
+    '''
+    calculate iou of two polygons
+    :param vertices1: vertices of the first polygon
+    :param vertices2: vertices of the second polygon
+    :return: the iou, the intersection area, the union area
+    '''
+    img1 = Image.new('L', (w, h), 0)
+    ImageDraw.Draw(img1).polygon(vertices1, outline=1, fill=1)
+    mask1 = np.array(img1)
+    img2 = Image.new('L', (w, h), 0)
+    ImageDraw.Draw(img2).polygon(vertices2, outline=1, fill=1)
+    mask2 = np.array(img2)
+    intersection = np.logical_and(mask1, mask2)
+    union = np.logical_or(mask1, mask2)
+    nu = np.sum(intersection)
+    de = np.sum(union)
+    if de!=0:
+        return nu*1.0/de, nu, de
+    else:
+        return 0, nu, de
+
+def label2vertex(labels):
+    '''
+    convert 1D labels to 2D vertices coordinates
+    :param labels: 1D labels
+    :return: 2D vertices coordinates: [(x1, y1),(x2,y2),...]
+    '''
+    vertices = []
+    for label in labels:
+        if (label == 784):
+            break
+        vertex = ((label % 28) * 8, (label / 28) * 8)
+        vertices.append(vertex)
+    return vertices
+
+def getbboxfromkps(kps,h,w):
+    '''
+
+    :param kps:
+    :return:
+    '''
+    min_c = np.min(np.array(kps), axis=0)
+    max_c = np.max(np.array(kps), axis=0)
+    object_h = max_c[1] - min_c[1]
+    object_w = max_c[0] - min_c[0]
+    h_extend = int(round(0.1 * object_h))
+    w_extend = int(round(0.1 * object_w))
+    min_row = np.maximum(0, min_c[1] - h_extend)
+    min_col = np.maximum(0, min_c[0] - w_extend)
+    max_row = np.minimum(h, max_c[1] + h_extend)
+    max_col = np.minimum(w, max_c[0] + w_extend)
+    return (min_row,min_col,max_row,max_col)
+
+def img2tensor(img):
+    '''
+
+    :param img:
+    :return:
+    '''
+    img = np.rollaxis(img,2,0)
+    return torch.from_numpy(img)
+
+def tensor2img(tensor):
+    '''
+
+    :param tensor:
+    :return:
+    '''
+    img = (tensor.numpy()*255).astype('uint8')
+    img = np.rollaxis(img,0,3)
+    return img