fix bug

data_generator.py

@@ -88,6 +88,7 @@ def data_gen(usage):
             batch_x[i_batch, :, :, 3] = trimap / 255.
             batch_y[i_batch, :, :, 0] = alpha / 255.
             # store trimap
+            trimap = np.array(trimap).astype(np.uint8)
             trimap_add(alpha / 255., trimap)
 
             i += 1

trimap_dict.py

@@ -1,14 +1,8 @@
-import cv2 as cv
-import numpy as np
-
-
 def trimap_init():
     global trimap_dict
     trimap_dict = dict()
-    global hit
-    hit = 0
-    global miss
-    miss = 0
+    global status
+    status = {'hit': 0, 'miss': 0}
 
 
 def trimap_add(alpha, trimap):
@@ -18,26 +12,22 @@ def trimap_add(alpha, trimap):
 
 
 def trimap_get(alpha):
-    global hit
-    global miss
     key = hash(str(alpha))
     if key in trimap_dict.keys():
-        hit += 1
+        status['hit'] = status['hit'] + 1
         return trimap_dict[key]
     else:
-        miss += 1
+        status['miss'] = status['miss'] + 1
         return None
 
 
 def trimap_clear(epoch):
-    global hit
-    global miss
+    global status
     size = len(trimap_dict)
     with open("training.txt", "a") as file:
-        file.write("Epoch %d, cleaning %d trimaps, hit=%d, miss=%d." % (epoch, size, hit, miss))
+        file.write("Epoch %d, cleaning %d trimaps, hit=%d, miss=%d.\n" % (epoch, size, status['hit'], status['miss']))
     trimap_dict.clear()
-    hit = 0
-    miss = 0
+    status = {'hit': 0, 'miss': 0}
 
 
 if __name__ == '__main__':