python-wml 3.0.0__py3-none-any.whl

wml/iotoolkit/aic_keypoint.py

@@ -0,0 +1,115 @@
+import json
+import numpy as np
+import os.path as osp
+import math
+
+'''
+       0: 'Right Shoulder',
+        1: 'Right Elbow',
+        2: 'Right Wrist',
+        3: 'Left Shoulder',
+        4: 'Left Elbow',
+        5: 'Left Wrist',
+        6: 'Right Hip',
+        7: 'Right Knee',
+        8: 'Right Ankle',
+        9: 'Left Hip',
+        10: 'Left Knee',
+        11: 'Left Ankle',
+        12: 'Head top',
+        13: 'Neck'
+'''
+def read_aic_keypoint(file_path,sub_dir_name=None):
+    with open(file_path,"r") as f:
+        datas = json.load(f)
+
+    res = []
+    for data in datas:
+        img_name = data['image_id']
+        kps = []
+        bboxes = []
+        human_data = data['human_annotations']
+        for k,v in data['keypoint_annotations'].items():
+            v = np.array(v,dtype=np.float32)
+            v = np.reshape(v,[-1,3])
+            mask = v[...,-1]==3
+            v[mask] = 0
+            bbox = np.array(human_data[k],dtype=np.float32)
+            kps.append(v)
+            bboxes.append(bbox)
+        kps = np.array(kps)
+        bboxes = np.array(bboxes)
+        if sub_dir_name is not None:
+            img_name = osp.join(sub_dir_name,img_name)
+        res.append([img_name+".jpg",bboxes,kps])
+
+    return res
+
+class Trans2COCO:
+    def __init__(self) -> None:
+        self.dst_idxs = [5,6,7,8,9,10,11,12,13,14,15,16]
+        self.src_idxs = [3,0,4,1,5,2,9,6,10,7,11,8]
+        self.coco_idxs = [0,1,2,3,4]
+    
+    @staticmethod
+    def get_head_pos(kp):
+        p0 = kp[12]
+        p1 = kp[13]
+        if p0[2]>0 and p1[2]>0:
+            dx = p0[0] - p1[0]
+            dy = p0[1] - p1[1]
+            dis = math.sqrt(dx*dx+dy*dy)*0.75
+            cx = (p0[0] + p1[0])/2
+            cy = (p0[1] + p1[1])/2
+            x0 = cx-dis
+            x1 = cx+dis
+            y0 = cy-dis
+            y1 = cy+dis
+            return [x0,y0,x1,y1]
+        else:
+            return None
+    
+    @staticmethod 
+    def kps_in_bbox(kps,bbox):
+        nr = kps.shape[0]
+        for i in range(nr):
+            if kps[i,2]<= 0:
+                return False
+            x = kps[i,0]
+            y = kps[i,1]
+            if x<bbox[0] or x>bbox[2] or y<bbox[1] or y>bbox[3]:
+                return False
+        return True
+
+
+    def __call__(self,mpii_kps,coco_kps=None):
+        if len(mpii_kps.shape)==2:
+            return self.trans_one(mpii_kps,coco_kps)
+        res = []
+        if coco_kps is not None:
+            for mp,coco in zip(mpii_kps,coco_kps):
+                res.append(self.trans_one(mp,coco))
+        else:
+            for mp in mpii_kps:
+                res.append(self.trans_one(mp))
+        return np.array(res)
+
+    def trans_one(self,mpii_kps,coco_kps=None):
+        '''
+        img: [RGB]
+        '''
+        res = np.zeros([17,3],dtype=np.float32)
+        res[self.dst_idxs] = mpii_kps[self.src_idxs]
+        if coco_kps is not None:
+            left_right_pairs = [[5,6],[11,12]]
+            is_good = True
+            for pair in left_right_pairs:
+                l,r = pair
+                if res[l,0]<res[r,0] or res[l,2]<0.1 or res[r,2]<0.1:
+                    is_good = False
+                    break
+
+            head_bbox = self.get_head_pos(mpii_kps)
+            if is_good and head_bbox is not None and self.kps_in_bbox(coco_kps[self.coco_idxs],head_bbox):
+                res[self.coco_idxs] = coco_kps[self.coco_idxs]
+        return res

wml/iotoolkit/baidu_mask_toolkit.py

@@ -0,0 +1,244 @@
+import json
+import cv2
+import wml.wml_utils as wmlu
+import numpy as np
+import os
+import cv2 as cv
+import sys
+import random
+from .common import *
+
+
+def get_files(dir_path, img_sub_dir=None,label_sub_dir=None):
+    if img_sub_dir is not None:
+        img_dir = os.path.join(dir_path, img_sub_dir)
+    else:
+        img_dir = dir_path
+
+    if label_sub_dir is not None:
+        label_dir = os.path.join(dir_path, label_sub_dir)
+    else:
+        label_dir = dir_path
+    res = []
+    json_files = wmlu.recurse_get_filepath_in_dir(label_dir,suffix=".json")
+    for jf in json_files:
+        base_name = wmlu.base_name(jf)
+        igf = os.path.join(img_dir, base_name + ".jpg")
+        if os.path.exists(igf):
+            res.append((igf, jf))
+        else:
+            print(f"ERROR: Find {igf} faild, json file is {jf}")
+
+    return res
+
+
+class BaiDuMaskData(object):
+    def __init__(self, trans_label=None,shuffle=False, img_sub_dir=None,label_sub_dir=None,ignored_labels=[],label_map={},
+                 allowed_labels_fn=None,overlap=True):
+        self.files = None
+        self.shuffle = shuffle
+        self.trans_label = trans_label
+        self.img_sub_dir = img_sub_dir
+        self.label_sub_dir = label_sub_dir
+        self.ignored_labels = ignored_labels
+        self.label_map = label_map
+        self.allowed_labels_fn = None if allowed_labels_fn is None or (isinstance(allowed_labels_fn,list ) and len(allowed_labels_fn)==0) else allowed_labels_fn
+        self.overlap = overlap
+        if self.allowed_labels_fn is not None and isinstance(self.allowed_labels_fn,list):
+            self.allowed_labels_fn = lambda x:x in allowed_labels_fn
+
+    def read_data(self, dir_path):
+        self.files = get_files(dir_path, self.img_sub_dir,self.label_sub_dir)
+        if self.shuffle:
+            random.shuffle(self.files)
+
+    def __len__(self):
+        return len(self.files)
+
+    def get_items(self,beg=0,end=None,filter=None):
+        '''
+        :return:
+        binary_masks [N,H,W], value is 0 or 1,
+        full_path,img_size,category_ids,category_names,boxes,binary_masks,area,is_crowd,num_annotations_skipped
+        '''
+        if end is None:
+            end = len(self.files)
+        if beg is None:
+            beg = 0
+        for i, (img_file, json_file) in enumerate(self.files[beg:end]):
+            if filter is not None and not filter(img_file,json_file):
+                continue
+            print(img_file,json_file)
+            sys.stdout.write('\r>> read data %d/%d' % (i + 1, len(self.files)))
+            sys.stdout.flush()
+            image, annotations_list = self.read_json(json_file,img_file)
+            masks = [ann["segmentation"] for ann in annotations_list]
+            labels = [ann['category_id'] for ann in annotations_list]
+            labels_names = [ann['category_name'] for ann in annotations_list]
+            bboxes = [ann['bbox'] for ann in annotations_list]
+            if len(masks) > 0:
+                try:
+                    masks = np.stack(masks, axis=0)
+                except:
+                    print("ERROR: stack masks faild.")
+                    masks = 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 = self.read_json(json_file,use_semantic=False)
+            labels = [ann['category_id'] for ann in annotations_list]
+            labels_names = [ann['category_name'] for ann in annotations_list]
+            bboxes = [ann['bbox'] for ann in annotations_list]
+            #file, img_size,category_ids, labels_text, bboxes, binary_mask, area, is_crowd, _
+            yield DetBboxesData(img_file, [image['height'], image['width']], labels, bboxes, None)
+
+    @staticmethod
+    def write_json(save_path,mask,labels2name,labels2color,label_trans=None,epsilon=None):
+        '''
+        save_path: json save path
+        mask:[H,W] or [N,H,W] uint_8, N为目标个数, 分辨率为原图分辨率
+        labels2name: list(tuple) (label,label_name}, len(labels2name)为目标个数
+        labels2color: {label,label_color}/list of colors, len(labels2color)=类别数
+        '''
+        shapes = []
+        for i,(k,n) in enumerate(labels2name):
+            if len(mask.shape)==2:
+                mask_tmp = (mask==k).astype(np.uint8)
+            else:
+                mask_tmp = mask[i]
+            if np.sum(mask_tmp)==0:
+                continue
+            contours, hierarchy = cv.findContours(mask_tmp, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
+            for cont in contours:
+                if epsilon is not None:
+                    cont = cv2.approxPolyDP(cont,epsilon,True)
+                points = cont
+                if len(cont.shape) == 3 and cont.shape[1] == 1:
+                    points = np.squeeze(points, axis=1)
+                points = points.tolist()
+                if len(points)<3:
+                    continue
+                tmp_data = {'name':n}
+                if label_trans is not None:
+                    tmp_data['labelIdx'] = int(label_trans[k])
+                else:
+                    tmp_data['labelIdx'] = int(k)
+                tmp_data["points"] = points
+                tmp_data['color'] = labels2color[k]
+                shapes.append(tmp_data)
+        with open(save_path,"w") as f:
+            json.dump(shapes,f)
+
+    def read_json(self,file_path,img_path,use_semantic=True):
+        annotations_list = []
+        image = {}
+        img = cv.imread(img_path)
+        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(img.shape[1])
+                img_height = int(img.shape[0])
+                image["height"] = int(img_height)
+                image["width"] = int(img_width)
+                image["file_name"] = wmlu.base_name(file_path)
+                for shape in json_data:
+                    label = shape["labelIdx"]
+                    label_name = shape['name']
+                    if self.ignored_labels is not None and label in self.ignored_labels:
+                        continue
+                    if self.allowed_labels_fn is not None and not self.allowed_labels_fn(label):
+                        continue
+                    if self.label_map is not None and label in self.label_map:
+                        label = self.label_map[label]
+                    mask = np.zeros(shape=[img_height, img_width], dtype=np.uint8)
+                    all_points = np.array([shape["points"]]).astype(np.int32)
+                    if len(all_points) < 1:
+                        continue
+                    points = np.transpose(all_points[0])
+                    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)
+                    if use_semantic:
+                        segmentation = cv.drawContours(mask, all_points, -1, color=(1), thickness=cv.FILLED)
+                        annotations_list.append({"bbox": (xmin, ymin, xmax - xmin + 1, ymax - ymin + 1),
+                                                 "segmentation": segmentation,
+                                                 "category_id": label,
+                                                 "category_name": label_name,
+                                                 "points_x": x,
+                                                 "points_y": y})
+                    else:
+                        annotations_list.append({"bbox": (xmin, ymin, xmax - xmin + 1, ymax - ymin + 1),
+                                             "category_id": label,
+                                             "category_name": label_name,
+                                             "points_x": x,
+                                             "points_y": y})
+            except:
+                print(f"Read file {os.path.basename(file_path)} faild.")
+                pass
+        if use_semantic and self.overlap:
+            '''
+            Each pixel only belong to one classes, and the latter annotation will overwrite the previous
+            '''
+            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
+
+
+if __name__ == "__main__":
+    id = 0
+    # data_statistics("/home/vghost/ai/mldata/qualitycontrol/rdatasv3")
+    import wml.img_utils as wmli
+    import wml.object_detection2.visualization as odv
+    import matplotlib.pyplot as plt
+
+    ID2NAME = {1:"person",2:"seat_belt"}
+
+
+    data = BaiDuMaskData( label_sub_dir="label",shuffle=False)
+    data.read_data("/home/wj/ai/mldata1/safety_belt/train_1")
+
+    def filter(x):
+        return x in ['general-single', 'parking', 'temporary', 'general-horizontal']
+        # 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]}"
+
+        if len(category_ids) == 0:
+            continue
+
+        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/base_dataset.py

@@ -0,0 +1,210 @@
+import random
+import os
+import os.path as osp
+import wml.wml_utils as wmlu
+from .common import resample,ignore_case_dict_label_text2id
+from abc import ABCMeta, abstractmethod
+from functools import partial
+
+
+
+class BaseDataset(metaclass=ABCMeta):
+    def __init__(self,label_text2id=None,
+                      filter_empty_files=False,filter_error=False,resample_parameters=None,shuffle=True,silent=False,keep_no_ann_imgs=False,absolute_coord=True,mask_on=False):
+        '''
+        filter_empty_files: remove files without any objects
+        keep_no_ann_imgs: keep images without annotation files
+        label_text2id: func(name)->int
+        '''
+        self.files = None
+        if isinstance(label_text2id,dict):
+            self.label_text2id = partial(ignore_case_dict_label_text2id,
+                    dict_data=wmlu.trans_dict_key2lower(label_text2id))
+        else:
+            self.label_text2id = label_text2id
+        self.filter_empty_files = filter_empty_files
+        self.filter_error = filter_error
+
+        if resample_parameters is not None:
+            self.resample_parameters = {}
+            for k,v in resample_parameters.items():
+                if isinstance(k,(str,bytes)) and self.label_text2id is not None:
+                    k = self.label_text2id(k)
+                self.resample_parameters[k] = v
+            print("resample parameters")
+            wmlu.show_dict(self.resample_parameters)
+        else:
+            self.resample_parameters = None
+
+        self.shuffle = shuffle
+        self.silent = silent
+        self.keep_no_ann_imgs = keep_no_ann_imgs
+        self.absolute_coord = absolute_coord
+        self.mask_on = mask_on
+        pass
+
+    def get_files_from_data_list_file(self,file_path):
+        res = []
+        with open(file_path,"r") as f:
+            lines = f.readlines()
+            for line in lines:
+                line = line.strip()
+                if "," in line:
+                    line = line.split(",")
+                res.append(line)
+        return res
+
+    def find_files(self,dir_path,img_suffix):
+        dir_path = self.process_path(dir_path)
+        if isinstance(dir_path,(list,tuple)) and len(dir_path)==1 and isinstance(dir_path[0],(str,bytes)):
+            dir_path = dir_path[0]
+
+        if isinstance(dir_path,str):
+            #从单一目录读取数据
+            print(f"Read {dir_path}")
+            if not os.path.exists(dir_path):
+                print(f"Data path {dir_path} not exists.")
+                return False
+            files = self.find_files_in_dir_or_dl(dir_path,img_suffix=img_suffix)
+        elif isinstance(dir_path,(list,tuple)):
+            if isinstance(dir_path[0],(str,bytes)) and os.path.isdir(dir_path[0]):
+                '''
+                从多个目录读取数据
+                Example:
+                [/dataset/dir0,/dataset/dir1,/dataset/dir2]
+                '''
+                files = self.find_files_in_dirs(dir_path, img_suffix=img_suffix)
+            elif isinstance(dir_path[0],(list,tuple)) and not isinstance(dir_path[0][1],(str,bytes)):
+                '''
+                用于对不同目录进行不同的重采样
+                Example:
+                [(/dataset/dir0,3),(/dataset/dir1,1),(/dataset/dir2,100)]
+                '''
+                files = self.find_files_in_dirs_with_repeat(dir_path,img_suffix=img_suffix)
+            else:
+                #文件列表
+                files = dir_path
+        
+        return files
+
+    def find_files_in_dir_or_dl(self,dir_path,img_suffix):
+        if osp.isfile(dir_path) and osp.splitext(dir_path)[-1].lower() == ".dl":
+            files = self.get_files_from_data_list_file(dir_path)
+        else:
+            files = self.find_files_in_dir(dir_path,img_suffix=img_suffix)
+
+        return files
+
+    @abstractmethod
+    def find_files_in_dir(self,dir_path,img_suffix):
+        pass
+
+    def find_files_in_dirs(self,dirs,img_suffix=".jpg"):
+        all_files = []
+        for dir_path in dirs:
+            files = self.find_files_in_dir_or_dl(dir_path,img_suffix)
+            print(f"Find {len(files)} in {dir_path}")
+            all_files.extend(files)
+
+        return all_files
+
+    def find_files_in_dirs_with_repeat(self,dirs,img_suffix=".jpg"):
+        all_files = []
+        raw_nr = 0
+        for dir_path,repeat_nr in dirs:
+            files = self.find_files_in_dir_or_dl(dir_path,img_suffix)
+            files_nr = len(files)
+            if repeat_nr>1:
+                files = list(files)*int(repeat_nr)
+
+            print(f"Find {files_nr} in {dir_path}, repeat to {len(files)} files")
+            raw_nr += files_nr
+            all_files.extend(list(files))
+
+        print(f"Total find {raw_nr} files, repeat to {len(all_files)} files.")
+        return all_files
+
+    def read_data(self,dir_path,img_suffix=".jpg;;.bmp;;.png;;.jpeg"):
+        _files = self.find_files(dir_path,img_suffix=img_suffix)
+        if self.filter_empty_files and self.label_text2id:
+            _files = self.apply_filter_empty_files(_files)
+        elif self.filter_error:
+            _files = self.apply_filter_error_files(_files)
+        if self.resample_parameters is not None and self.label_text2id:
+            _files = self.resample(_files)
+        
+        self.files = _files
+
+        if isinstance(dir_path,(str,bytes)):
+            print(f"Total find {len(self.files)} in {dir_path}")
+        else:
+            print(f"Total find {len(self.files)}")
+
+        if self.shuffle:
+            random.shuffle(self.files)
+
+        print("Files")
+        wmlu.show_list(self.files[:100])
+        if len(self.files)>100:
+            print("...")
+
+    def apply_filter_empty_files(self,files):
+        new_files = []
+        for fs in files:
+            try:
+                labels,labels_names = self.get_labels(fs)
+                is_none = [x is None for x in labels]
+                if not all(is_none):
+                    new_files.append(fs)
+                else:
+                    print(f"File {fs[1]} is empty, remove from dataset, labels names {labels_names}, labels {labels}")
+            except Exception as e:
+                print(f"Read {fs[1]} faild, info: {e}.")
+                pass
+
+        return new_files
+
+    def apply_filter_error_files(self,files):
+        new_files = []
+        for fs in files:
+            try:
+                labels,labels_names = self.get_labels(fs)
+                new_files.append(fs)
+            except Exception as e:
+                print(f"Read {fs[1]} faild, info: {e}.")
+                pass
+
+        return new_files
+
+    @abstractmethod
+    def get_labels(self,fs):
+        #return labels,labels_names
+        pass
+    
+    def resample(self,files):
+        all_labels = []
+        for fs in files:
+            try:
+                labels,_ = self.get_labels(fs)
+                all_labels.append(labels)
+            except Exception as e:
+                print(f"Basedataset resample error: {e}")
+
+        return resample(files,all_labels,self.resample_parameters)
+
+    def __len__(self):
+        return len(self.files)
+
+    @staticmethod
+    def process_path(path):
+        if isinstance(path,(list,tuple)):
+            return [BaseDataset.process_path(x) for x in path]
+        if isinstance(path,(int,float)):
+            return path
+        if isinstance(path,(str,bytes)):
+            return osp.abspath(osp.expanduser(path))
+        return path
+
+    
+    
+