add

dataset/casia_hwdb.py

@@ -0,0 +1,87 @@
+"""
+
+this is a wrapper handle CASIA_HWDB dataset
+since original data is complicated
+we using this class to get .png and label from raw
+.gnt data
+
+"""
+import struct
+import numpy as np
+import cv2
+
+
+class CASIAHWDBGNT(object):
+    """
+    A .gnt file may contains many images and charactors
+    """
+
+    def __init__(self, f_p):
+        self.f_p = f_p
+
+    def get_data_iter(self):
+        header_size = 10
+        with open(self.f_p, 'rb') as f:
+            while True:
+                header = np.fromfile(f, dtype='uint8', count=header_size)
+                if not header.size: 
+                    break
+                sample_size = header[0] + (header[1]<<8) + (header[2]<<16) + (header[3]<<24)
+                tagcode = header[5] + (header[4]<<8)
+                width = header[6] + (header[7]<<8)
+                height = header[8] + (header[9]<<8)
+                if header_size + width*height != sample_size:
+                    break
+                image = np.fromfile(f, dtype='uint8', count=width*height).reshape((height, width))
+                yield image, tagcode
+
+
+def resize_padding_or_crop(target_size, ori_img, padding_value=255):
+    if len(ori_img.shape) == 3:
+        res = np.zeros([ori_img.shape[0], target_size, target_size])
+    else:
+        res = np.ones([target_size, target_size])*padding_value
+        end_x = target_size
+        end_y = target_size
+        start_x = 0
+        start_y = 0
+        if ori_img.shape[0] < target_size:
+            end_x = int((target_size + ori_img.shape[0])/2)
+        if ori_img.shape[1] < target_size:
+            end_y = int((target_size + ori_img.shape[1])/2)
+        if ori_img.shape[0] < target_size:
+            start_x = int((target_size - ori_img.shape[0])/2)
+        if ori_img.shape[1] < target_size:
+            start_y = int((target_size - ori_img.shape[1])/2)
+        res[start_x:end_x, start_y:end_y] = ori_img
+        return np.array(res, dtype=np.uint8)
+
+if __name__ == "__main__":
+    gnt = CASIAHWDBGNT('samples/1001-f.gnt')
+
+    full_img = np.zeros([800, 800], dtype=np.uint8)
+    charset = []
+    i = 0
+    for img, tagcode in gnt.get_data_iter():
+        cv2.imshow('rr', img)
+        
+        try:
+            label = struct.pack('>H', tagcode).decode('gb2312')
+            cv2.waitKey(0)
+            print(label)
+            # img_padded = resize_padding_or_crop(80, img)
+            # col_idx = i%10
+            # row_idx = i//10
+            # full_img[row_idx*80:(row_idx+1)*80, col_idx*80:(col_idx+1)*80] = img_padded
+            # charset.append(label.replace('\x00', ''))
+            # if i >= 99:
+            #     cv2.imshow('rrr', full_img)
+            #     cv2.imwrite('sample.png', full_img)
+            #     cv2.waitKey(0)
+            #     print(charset)
+            #     break
+            # i += 1
+        except Exception as e:
+            # print(e.with_traceback(0))
+            print('decode error')
+            continue