python-wml 3.0.0__py3-none-any.whl

wml/iotoolkit/cityscapes_toolkit.py

@@ -0,0 +1,183 @@
+import json
+import wml.wml_utils as wmlu
+import numpy as np
+import os
+import cv2 as cv
+import sys
+import random
+from wml.iotoolkit.labelme_toolkit import get_labels_and_bboxes
+from .common import *
+
+def read_json(file_path):
+    annotations_list = []
+    image = {}
+    with open(file_path, "r", encoding="gb18030") as f:
+        print(file_path)
+        data_str = f.read()
+        try:
+            json_data = json.loads(data_str)
+            img_width = int(json_data["imgWidth"])
+            img_height = int(json_data["imgHeight"])
+            image["height"] = int(img_height)
+            image["width"] = int(img_width)
+            image["file_name"] = wmlu.base_name(file_path)
+            for shape in json_data["objects"]:
+                mask = np.zeros(shape=[img_height, img_width], dtype=np.uint8)
+                all_points = np.array([shape["polygon"]]).astype(np.int32)
+                if len(all_points) < 1:
+                    continue
+                points = np.transpose(all_points[0])
+                print(points.shape)
+                x, y = np.vsplit(points, 2)
+                x = np.reshape(x, [-1])
+                y = np.reshape(y, [-1])
+                x = np.minimum(np.maximum(0, x), img_width - 1)
+                y = np.minimum(np.maximum(0, y), img_height - 1)
+                xmin = np.min(x)
+                xmax = np.max(x)
+                ymin = np.min(y)
+                ymax = np.max(y)
+                segmentation = cv.drawContours(mask, all_points, -1, color=(1), thickness=cv.FILLED)
+                label = shape["label"]
+                annotations_list.append({"bbox": (xmin, ymin, xmax - xmin + 1, ymax - ymin + 1),
+                                         "segmentation": segmentation,
+                                         "category_id": label,
+                                         "points_x": x,
+                                         "points_y": y})
+        except:
+            print(f"Read file {os.path.basename(file_path)} faild.")
+            pass
+
+    if len(annotations_list) > 2:
+        mask = 1 - annotations_list[-1]['segmentation']
+        for i in reversed(range(len(annotations_list) - 1)):
+            annotations_list[i]['segmentation'] = np.logical_and(annotations_list[i]['segmentation'], mask)
+            mask = np.logical_and(mask, 1 - annotations_list[i]['segmentation'])
+
+    return image, annotations_list
+
+def get_files(dir_path, sub_dir_name):
+    img_dir = os.path.join(dir_path, 'leftImg8bit', sub_dir_name)
+    label_dir = os.path.join(dir_path, 'gtFine', sub_dir_name)
+    dir_names = wmlu.get_subdir_in_dir(label_dir)
+    res = []
+    for dir_name in dir_names:
+        l_label_dir = os.path.join(label_dir, dir_name)
+        l_img_dir = os.path.join(img_dir, dir_name)
+        json_files = wmlu.recurse_get_filepath_in_dir(l_label_dir,suffix=".json")
+        for jf in json_files:
+            base_name = wmlu.base_name(jf)[:-15]+"leftImg8bit"
+            igf = os.path.join(l_img_dir, base_name + ".png")
+            if os.path.exists(igf):
+                res.append((igf, jf))
+            else:
+                print(f"ERROR: Find {igf} faild, json file is {jf}")
+
+    return res
+
+class CityscapesData(object):
+    def __init__(self, label_text2id=None, shuffle=False, sub_dir_name="train"):
+        self.files = None
+        self.label_text2id = label_text2id
+        self.shuffle = shuffle
+        self.sub_dir_name = sub_dir_name
+
+    def read_data(self, dir_path):
+        self.files = get_files(dir_path, self.sub_dir_name)
+        if self.shuffle:
+            random.shuffle(self.files)
+
+    def get_items(self):
+        '''
+        :return:
+        full_path,img_size,category_ids,category_names,boxes,binary_masks,area,is_crowd,num_annotations_skipped
+        '''
+        for i, (img_file, json_file) in enumerate(self.files):
+            print(img_file,json_file)
+            sys.stdout.write('\r>> read data %d/%d' % (i + 1, len(self.files)))
+            sys.stdout.flush()
+            image, annotations_list = read_json(json_file)
+            labels_names, bboxes = get_labels_and_bboxes(image, annotations_list)
+            masks = [ann["segmentation"] for ann in annotations_list]
+            if len(masks) > 0:
+                try:
+                    masks = np.stack(masks, axis=0)
+                except:
+                    print("ERROR: stack masks faild.")
+                    masks = None
+
+            if self.label_text2id is not None:
+                labels = [self.label_text2id(x) for x in labels_names]
+            else:
+                labels = None
+
+            yield DetData(img_file, [image['height'], image['width']], labels, labels_names, bboxes, masks, None, None, None)
+
+    def get_boxes_items(self):
+        '''
+        :return:
+        full_path,img_size,category_ids,boxes,is_crowd
+        '''
+        for i, (img_file, json_file) in enumerate(self.files):
+            sys.stdout.write('\r>> read data %d/%d' % (i + 1, len(self.files)))
+            sys.stdout.flush()
+            image, annotations_list = read_json(json_file)
+            labels_names, bboxes = get_labels_and_bboxes(image, annotations_list)
+            labels = [self.label_text2id(x) for x in labels_names]
+            yield DetBboxesData(img_file, [image['height'], image['width']], labels,bboxes, None)
+
+if __name__ == "__main__":
+    # data_statistics("/home/vghost/ai/mldata/qualitycontrol/rdatasv3")
+    import wml.img_utils as wmli
+    import object_detection_tools.visualization as odv
+    import matplotlib.pyplot as plt
+
+    NAME2ID = {}
+    ID2NAME = {}
+    id = 0
+
+    def name_to_id(x):
+        global id
+        if x in NAME2ID:
+            return NAME2ID[x]
+        else:
+            NAME2ID[x] = id
+            ID2NAME[id] = x
+            id += 1
+            return NAME2ID[x]
+
+
+    data = CityscapesData(label_text2id=name_to_id, shuffle=False)
+    # data.read_data("/data/mldata/qualitycontrol/rdatasv5_splited/rdatasv5")
+    # data.read_data("/home/vghost/ai/mldata2/qualitycontrol/rdatav10_preproc")
+    # data.read_data("/home/vghost/ai/mldata2/qualitycontrol/rdatasv10_neg_preproc")
+    data.read_data(wmlu.home_dir("ai/mldata/cityscapes"))
+    def filter(x):
+        return x in ['road','parking']
+        #return x in ['person', 'parking']
+        #return x in ['terrain']
+        #return x in ['car']
+    # data.read_data("/home/vghost/ai/mldata2/qualitycontrol/x")
+    for x in data.get_items():
+        full_path, img_info, category_ids, category_names, boxes, binary_mask, area, is_crowd, num_annotations_skipped = x
+        img = wmli.imread(full_path)
+
+        def text_fn(classes, scores):
+            return f"{ID2NAME[classes]}"
+
+        is_keep = [filter(x) for x in category_names]
+        category_ids = np.array(category_ids)[is_keep]
+        boxes = np.array(boxes)[is_keep]
+        binary_mask = np.array(binary_mask)[is_keep]
+        if len(category_ids)==0:
+            continue
+
+        wmlu.show_dict(NAME2ID)
+        odv.draw_bboxes_and_maskv2(
+            img=img, classes=category_ids, scores=None, bboxes=boxes, masks=binary_mask, color_fn=None,
+            text_fn=text_fn, thickness=4,
+            show_text=True,
+            fontScale=0.8)
+        plt.figure()
+        plt.imshow(img)
+        plt.show()

wml/iotoolkit/classification_data_statistics.py

@@ -0,0 +1,25 @@
+import wml.wml_utils as wmlu
+
+def labels_statistics(labels,name="",classes_names=None):
+    counter = wmlu.Counter()
+    for l in labels:
+        if isinstance(l,(list,tuple)):
+            l = l[0]
+        counter.add(l)
+    
+    
+    counter = list(counter.items())
+    counter.sort(key=lambda x:x[1],reverse=True)
+    classes = []
+    for l,v in counter:
+        classes.append(l)
+    print(f"Total {len(classes)} classes")
+    print(list(classes))
+    if len(name)>0:
+        print(name)
+    for l,v in counter:
+        if classes_names is not None:
+            l = classes_names[l]
+        print(f"{l:>32}: {v:<8},  {v*100.0/len(labels):>3.2f}%")
+
+    print(f"Total {len(labels)} samples.")

wml/iotoolkit/coco_data_fwd.py

@@ -0,0 +1,225 @@
+import numpy as np
+
+COCO_KP_NR = 17
+ID_TO_COMPRESSED_ID = {}  #第一个类别的编号为1
+ZERO_ID_TO_COMPRESSED_ID = {}  #第一个类别的编号为0
+COMPRESSED_ID_TO_ID = {} #第一个类别的编号为1
+ZERO_COMPRESSED_ID_TO_ID = {} #第一个类别的编号为0
+
+ID_TO_TEXT={1: {u'supercategory': u'person', u'id': 1, u'name': u'person'},
+            2: {u'supercategory': u'vehicle', u'id': 2, u'name': u'bicycle'},
+            3: {u'supercategory': u'vehicle', u'id': 3, u'name': u'car'},
+            4: {u'supercategory': u'vehicle', u'id': 4, u'name': u'motorcycle'},
+            5: {u'supercategory': u'vehicle', u'id': 5, u'name': u'airplane'},
+            6: {u'supercategory': u'vehicle', u'id': 6, u'name': u'bus'},
+            7: {u'supercategory': u'vehicle', u'id': 7, u'name': u'train'},
+            8: {u'supercategory': u'vehicle', u'id': 8, u'name': u'truck'},
+            9: {u'supercategory': u'vehicle', u'id': 9, u'name': u'boat'},
+            10: {u'supercategory': u'outdoor', u'id': 10, u'name': u'traffic light'},
+            11: {u'supercategory': u'outdoor', u'id': 11, u'name': u'fire hydrant'},
+            13: {u'supercategory': u'outdoor', u'id': 13, u'name': u'stop sign'},
+            14: {u'supercategory': u'outdoor', u'id': 14, u'name': u'parking meter'},
+            15: {u'supercategory': u'outdoor', u'id': 15, u'name': u'bench'},
+            16: {u'supercategory': u'animal', u'id': 16, u'name': u'bird'},
+            17: {u'supercategory': u'animal', u'id': 17, u'name': u'cat'},
+            18: {u'supercategory': u'animal', u'id': 18, u'name': u'dog'},
+            19: {u'supercategory': u'animal', u'id': 19, u'name': u'horse'},
+            20: {u'supercategory': u'animal', u'id': 20, u'name': u'sheep'},
+            21: {u'supercategory': u'animal', u'id': 21, u'name': u'cow'},
+            22: {u'supercategory': u'animal', u'id': 22, u'name': u'elephant'},
+            23: {u'supercategory': u'animal', u'id': 23, u'name': u'bear'},
+            24: {u'supercategory': u'animal', u'id': 24, u'name': u'zebra'},
+            25: {u'supercategory': u'animal', u'id': 25, u'name': u'giraffe'},
+            27: {u'supercategory': u'accessory', u'id': 27, u'name': u'backpack'},
+            28: {u'supercategory': u'accessory', u'id': 28, u'name': u'umbrella'},
+            31: {u'supercategory': u'accessory', u'id': 31, u'name': u'handbag'},
+            32: {u'supercategory': u'accessory', u'id': 32, u'name': u'tie'},
+            33: {u'supercategory': u'accessory', u'id': 33, u'name': u'suitcase'},
+            34: {u'supercategory': u'sports', u'id': 34, u'name': u'frisbee'},
+            35: {u'supercategory': u'sports', u'id': 35, u'name': u'skis'},
+            36: {u'supercategory': u'sports', u'id': 36, u'name': u'snowboard'},
+            37: {u'supercategory': u'sports', u'id': 37, u'name': u'sports ball'},
+            38: {u'supercategory': u'sports', u'id': 38, u'name': u'kite'},
+            39: {u'supercategory': u'sports', u'id': 39, u'name': u'baseball bat'},
+            40: {u'supercategory': u'sports', u'id': 40, u'name': u'baseball glove'},
+            41: {u'supercategory': u'sports', u'id': 41, u'name': u'skateboard'},
+            42: {u'supercategory': u'sports', u'id': 42, u'name': u'surfboard'},
+            43: {u'supercategory': u'sports', u'id': 43, u'name': u'tennis racket'},
+            44: {u'supercategory': u'kitchen', u'id': 44, u'name': u'bottle'},
+            46: {u'supercategory': u'kitchen', u'id': 46, u'name': u'wine glass'},
+            47: {u'supercategory': u'kitchen', u'id': 47, u'name': u'cup'},
+            48: {u'supercategory': u'kitchen', u'id': 48, u'name': u'fork'},
+            49: {u'supercategory': u'kitchen', u'id': 49, u'name': u'knife'},
+            50: {u'supercategory': u'kitchen', u'id': 50, u'name': u'spoon'},
+            51: {u'supercategory': u'kitchen', u'id': 51, u'name': u'bowl'},
+            52: {u'supercategory': u'food', u'id': 52, u'name': u'banana'},
+            53: {u'supercategory': u'food', u'id': 53, u'name': u'apple'},
+            54: {u'supercategory': u'food', u'id': 54, u'name': u'sandwich'},
+            55: {u'supercategory': u'food', u'id': 55, u'name': u'orange'},
+            56: {u'supercategory': u'food', u'id': 56, u'name': u'broccoli'},
+            57: {u'supercategory': u'food', u'id': 57, u'name': u'carrot'},
+            58: {u'supercategory': u'food', u'id': 58, u'name': u'hot dog'},
+            59: {u'supercategory': u'food', u'id': 59, u'name': u'pizza'},
+            60: {u'supercategory': u'food', u'id': 60, u'name': u'donut'},
+            61: {u'supercategory': u'food', u'id': 61, u'name': u'cake'},
+            62: {u'supercategory': u'furniture', u'id': 62, u'name': u'chair'},
+            63: {u'supercategory': u'furniture', u'id': 63, u'name': u'couch'},
+            64: {u'supercategory': u'furniture', u'id': 64, u'name': u'potted plant'},
+            65: {u'supercategory': u'furniture', u'id': 65, u'name': u'bed'},
+            67: {u'supercategory': u'furniture', u'id': 67, u'name': u'dining table'},
+            70: {u'supercategory': u'furniture', u'id': 70, u'name': u'toilet'},
+            72: {u'supercategory': u'electronic', u'id': 72, u'name': u'tv'},
+            73: {u'supercategory': u'electronic', u'id': 73, u'name': u'laptop'},
+            74: {u'supercategory': u'electronic', u'id': 74, u'name': u'mouse'},
+            75: {u'supercategory': u'electronic', u'id': 75, u'name': u'remote'},
+            76: {u'supercategory': u'electronic', u'id': 76, u'name': u'keyboard'},
+            77: {u'supercategory': u'electronic', u'id': 77, u'name': u'cell phone'},
+            78: {u'supercategory': u'appliance', u'id': 78, u'name': u'microwave'},
+            79: {u'supercategory': u'appliance', u'id': 79, u'name': u'oven'},
+            80: {u'supercategory': u'appliance', u'id': 80, u'name': u'toaster'},
+            81: {u'supercategory': u'appliance', u'id': 81, u'name': u'sink'},
+            82: {u'supercategory': u'appliance', u'id': 82, u'name': u'refrigerator'},
+            84: {u'supercategory': u'indoor', u'id': 84, u'name': u'book'},
+            85: {u'supercategory': u'indoor', u'id': 85, u'name': u'clock'},
+            86: {u'supercategory': u'indoor', u'id': 86, u'name': u'vase'},
+            87: {u'supercategory': u'indoor', u'id': 87, u'name': u'scissors'},
+            88: {u'supercategory': u'indoor', u'id': 88, u'name': u'teddy bear'},
+            89: {u'supercategory': u'indoor', u'id': 89, u'name': u'hair drier'},
+            90: {u'supercategory': u'indoor', u'id': 90, u'name': u'toothbrush'}}
+COMPRESSED_ID_TO_TEXT = {}
+
+
+
+KEYPOINTS_NAME = ["nose", #0
+                "left_eye", #1
+                "right_eye",#2
+                "left_ear",#3
+                "right_ear",#4
+                "left_shoulder",#5
+                "right_shoulder",#6
+                "left_elbow",#7
+                "right_elbow",#8
+                "left_wrist",#9
+                "right_wrist",#10
+                "left_hip",#11
+                "right_hip",#12
+                "left_knee",#13
+                "right_knee",#14
+                "left_ankle",#15
+                "right_ankle"]#16
+JOINTS_PAIR = [[0 , 1], [1 , 2], [2 , 0], [1 , 3], [2 , 4], [3 , 5], [4 , 6], [5 , 6], [5 , 11],
+[6 , 12], [11 , 12], [5 , 7], [7 , 9], [6 , 8], [8 , 10], [11 , 13], [13 , 15], [12 , 14], [14 , 16]]
+j = 1
+for i in range(1,81):
+    while j not in ID_TO_TEXT:
+        j += 1
+    ID_TO_COMPRESSED_ID[j] = i
+    ZERO_ID_TO_COMPRESSED_ID[j] = i-1
+    COMPRESSED_ID_TO_ID[i] = j
+    ZERO_COMPRESSED_ID_TO_ID[i-1] = j
+    j += 1
+
+for k,v in COMPRESSED_ID_TO_ID.items():
+    COMPRESSED_ID_TO_TEXT[k] = ID_TO_TEXT[v]['name']
+
+COCO_CLASSES_FREQ = {
+"person": 30.52,
+"car": 5.10,
+"chair": 4.48,
+"book": 2.87,
+"bottle": 2.83,
+"cup": 2.40,
+"dining table": 1.83,
+"bowl": 1.67,
+"traffic light": 1.50,
+"handbag": 1.44,
+"umbrella": 1.33,
+"bird": 1.26,
+"boat": 1.25,
+"truck": 1.16,
+"bench": 1.14,
+"sheep": 1.11,
+"banana": 1.10,
+"kite": 1.06,
+"motorcycle": 1.01,
+"backpack": 1.01,
+"potted plant": 1.01,
+"cow": 0.95,
+"wine glass": 0.92,
+"carrot": 0.91,
+"knife": 0.90,
+"broccoli": 0.85,
+"donut": 0.83,
+"bicycle": 0.83,
+"skis": 0.77,
+"vase": 0.77,
+"horse": 0.77,
+"tie": 0.76,
+"cell phone": 0.75,
+"orange": 0.74,
+"cake": 0.74,
+"sports ball": 0.74,
+"clock": 0.74,
+"suitcase": 0.72,
+"spoon": 0.72,
+"surfboard": 0.71,
+"bus": 0.71,
+"apple": 0.68,
+"pizza": 0.68,
+"tv": 0.68,
+"couch": 0.67,
+"remote": 0.66,
+"sink": 0.65,
+"skateboard": 0.64,
+"elephant": 0.64,
+"dog": 0.64,
+"fork": 0.64,
+"zebra": 0.62,
+"airplane": 0.60,
+"giraffe": 0.60,
+"laptop": 0.58,
+"tennis racket": 0.56,
+"teddy bear": 0.56,
+"cat": 0.55,
+"train": 0.53,
+"sandwich": 0.51,
+"bed": 0.49,
+"toilet": 0.48,
+"baseball glove": 0.44,
+"oven": 0.39,
+"baseball bat": 0.38,
+"hot dog": 0.34,
+"keyboard": 0.33,
+"snowboard": 0.31,
+"frisbee": 0.31,
+"refrigerator": 0.31,
+"mouse": 0.26,
+"stop sign": 0.23,
+"toothbrush": 0.23,
+"fire hydrant": 0.22,
+"microwave": 0.19,
+"scissors": 0.17,
+"bear": 0.15,
+"parking meter": 0.15,
+"toaster": 0.03,
+"hair drier": 0.02
+}
+ID_TO_FREQ = np.ones(shape=[91],dtype=np.float32)
+for i in range(1,91):
+    if i in ID_TO_TEXT:
+        name = ID_TO_TEXT[i]['name']
+        if name not in COCO_CLASSES_FREQ:
+            print(f"Error {i}/{name}")
+            raise ValueError(name)
+        else:
+            ID_TO_FREQ[i] = COCO_CLASSES_FREQ[name]
+COMPRESSED_ID_TO_FREQ = np.ones(shape=[81],dtype=np.float32)
+for i in range(1,81):
+    id = COMPRESSED_ID_TO_ID[i]
+    COMPRESSED_ID_TO_FREQ[i] = ID_TO_FREQ[id]
+
+'''
+APsmall% AP for small objects: area < 32^2
+APmedium% AP for medium objects: 32^2 < area < 96^2 
+APlarge% AP for large objects: area > 96^2
+'''

wml/iotoolkit/coco_keypoints.py

@@ -0,0 +1,118 @@
+from pycocotools.coco import COCO
+import numpy as np
+import copy
+
+class COCOKeypoints(object):
+    def __init__(self):
+        self.pixel_std = 200
+        self.num_joints = 17
+        pass
+
+    def read_data(self,file_path):
+        self.coco = COCO(file_path)
+        self.image_set_index = self.coco.getImgIds()
+        self.gt_db = self._load_coco_keypoint_annotations()
+        return self.gt_db
+
+    def image_path_from_index(self, index):
+        file_name = '%012d.jpg' % index
+        return file_name
+
+    def _load_coco_keypoint_annotations(self):
+        """ ground truth bbox and keypoints """
+        gt_db = []
+        for index in self.image_set_index:
+            gt_db.extend(self._load_coco_keypoint_annotation_kernal(index))
+        return gt_db
+
+
+    def _load_coco_keypoint_annotation_kernal(self, index):
+        """
+        coco ann: [u'segmentation', u'area', u'iscrowd', u'image_id', u'bbox', u'category_id', u'id']
+        iscrowd:
+            crowd instances are handled by marking their overlaps with all categories to -1
+            and later excluded in training
+        bbox:
+            [x1, y1, w, h]
+        :param index: coco image id
+        :return: db entry
+        """
+        im_ann = self.coco.loadImgs(index)[0]
+        width = im_ann['width']
+        height = im_ann['height']
+
+        annIds = self.coco.getAnnIds(imgIds=index, iscrowd=False)
+        objs = self.coco.loadAnns(annIds)
+
+        # sanitize bboxes
+        valid_objs = []
+        for obj in objs:
+            x, y, w, h = obj['bbox']
+            x1 = np.max((0, x))
+            y1 = np.max((0, y))
+            x2 = np.min((width - 1, x1 + np.max((0, w - 1))))
+            y2 = np.min((height - 1, y1 + np.max((0, h - 1))))
+            if obj['area'] > 0 and x2 >= x1 and y2 >= y1:
+                obj['clean_bbox'] = [x1, y1, x2 - x1, y2 - y1]
+                valid_objs.append(obj)
+        objs = valid_objs
+
+        rec = []
+        for obj in objs:
+            '''cls = self._coco_ind_to_class_ind[obj['category_id']]
+            if cls != 1:
+                continue'''
+            # ignore objs without keypoints annotation
+            if max(obj['keypoints']) == 0:
+                continue
+
+            joints_3d = np.zeros((self.num_joints, 3), dtype=np.float32)
+            joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float32)
+            for ipt in range(self.num_joints):
+                joints_3d[ipt, 0] = obj['keypoints'][ipt * 3 + 0]
+                joints_3d[ipt, 1] = obj['keypoints'][ipt * 3 + 1]
+                t_vis = obj['keypoints'][ipt * 3 + 2]
+                if t_vis > 1:
+                    t_vis = 1
+                joints_3d[ipt, 2] = t_vis
+                joints_3d_vis[ipt, 0] = t_vis
+                joints_3d_vis[ipt, 1] = t_vis
+                joints_3d_vis[ipt, 2] = 0
+
+            clean_bbox = copy.deepcopy(obj['clean_bbox'])
+            center, scale = self._box2cs(clean_bbox[:4])
+            rec.append({
+                'image_id':index,
+                'image': self.image_path_from_index(index),
+                'center': center,
+                'scale': scale,
+                'clean_bbox':clean_bbox,
+                'keypoints': joints_3d,
+                'joints_3d_vis': joints_3d_vis,
+                'filename': '',
+                'imgnum': 0,
+            })
+
+        return rec
+
+
+    def _box2cs(self, box):
+        x, y, w, h = box[:4]
+        return self._xywh2cs(x, y, w, h)
+
+    def _xywh2cs(self, x, y, w, h):
+        center = np.zeros((2), dtype=np.float32)
+        center[0] = x + w * 0.5
+        center[1] = y + h * 0.5
+
+        '''if w > self.aspect_ratio * h:
+            h = w * 1.0 / self.aspect_ratio
+        elif w < self.aspect_ratio * h:
+            w = h * self.aspect_ratio'''
+        scale = np.array(
+            [w * 1.0 / self.pixel_std, h * 1.0 / self.pixel_std],
+            dtype=np.float32)
+        if center[0] != -1:
+            scale = scale * 1.25
+
+        return center, scale

wml/iotoolkit/coco_keypoints_fmt2.py

@@ -0,0 +1,103 @@
+from pycocotools.coco import COCO
+import numpy as np
+import copy
+
+'''
+返回image层面的结果
+'''
+class COCOKeypointsFmt2(object):
+    def __init__(self):
+        self.pixel_std = 200
+        self.num_joints = 17
+        pass
+
+    def read_data(self,file_path):
+        self.coco = COCO(file_path)
+        self.image_set_index = self.coco.getImgIds()
+        self.gt_db = self._load_coco_keypoint_annotations()
+        return self.gt_db
+
+    def image_path_from_index(self, index):
+        file_name = '%012d.jpg' % index
+        return file_name
+
+    def _load_coco_keypoint_annotations(self):
+        """ ground truth bbox and keypoints """
+        gt_db = []
+        for index in self.image_set_index:
+            res = self._load_coco_keypoint_annotation_kernal(index)
+            if res is not None:
+                gt_db.append(res)
+        return gt_db
+
+
+    def _load_coco_keypoint_annotation_kernal(self, index):
+        """
+        coco ann: [u'segmentation', u'area', u'iscrowd', u'image_id', u'bbox', u'category_id', u'id']
+        iscrowd:
+            crowd instances are handled by marking their overlaps with all categories to -1
+            and later excluded in training
+        bbox:
+            [x1, y1, w, h]
+        :param index: coco image id
+        :return: db entry
+        """
+        im_ann = self.coco.loadImgs(index)[0]
+        width = im_ann['width']
+        height = im_ann['height']
+
+        annIds = self.coco.getAnnIds(imgIds=index, iscrowd=False)
+        objs = self.coco.loadAnns(annIds)
+
+        # sanitize bboxes
+        valid_objs = []
+        for obj in objs:
+            x, y, w, h = obj['bbox']
+            x1 = np.max((0, x))
+            y1 = np.max((0, y))
+            x2 = np.min((width - 1, x1 + np.max((0, w - 1))))
+            y2 = np.min((height - 1, y1 + np.max((0, h - 1))))
+            if obj['area'] > 0 and x2 >= x1 and y2 >= y1:
+                obj['clean_bbox'] = [x1, y1, x2 - x1, y2 - y1]
+                valid_objs.append(obj)
+        objs = valid_objs
+
+        if len(objs) == 0:
+            return None
+
+        res = {}
+        res['image'] = self.image_path_from_index(index)
+        all_kps = []
+        all_bboxes = []
+        all_areas = []
+        all_iscrowd = []
+        for obj in objs:
+            '''cls = self._coco_ind_to_class_ind[obj['category_id']]
+            if cls != 1:
+                continue'''
+            # ignore objs without keypoints annotation
+            if max(obj['keypoints']) == 0:
+                continue
+
+            joints_3d = np.zeros((self.num_joints, 3), dtype=np.float32)
+            for ipt in range(self.num_joints):
+                joints_3d[ipt, 0] = obj['keypoints'][ipt * 3 + 0]
+                joints_3d[ipt, 1] = obj['keypoints'][ipt * 3 + 1]
+                t_vis = obj['keypoints'][ipt * 3 + 2]
+                if t_vis > 1:
+                    t_vis = 1
+                joints_3d[ipt, 2] = t_vis
+
+            clean_bbox = np.array(obj['clean_bbox'])
+            bbox = clean_bbox
+            bbox[2:] = bbox[:2]+bbox[2:]
+            all_kps.append(joints_3d)
+            all_bboxes.append(bbox)
+            all_areas.append(obj['area'])
+            all_iscrowd.append(obj['iscrowd'])
+
+        res["keypoints"] = np.array(all_kps,dtype=np.float32)
+        res['bboxes'] = np.array(all_bboxes,dtype=np.float32)
+        res['area'] = np.array(all_areas,dtype=np.float32)
+        res['iscrowd'] = np.array(all_iscrowd,dtype=np.float32)
+        return res