naver/PUMP

1# Copyright 2022-present NAVER Corp.2# CC BY-NC-SA 4.03# Available only for non-commercial use4 5from pdb import set_trace as bb6import numpy as np7from PIL import Image, ImageOps, ImageEnhance8 9 10def grab( data, *fields ):11    ''' Called to extract fields from a dictionary12    '''13    if isinstance(data, dict):14        res = []15        for f in fields:16            res.append( data[f] )17        return res[0] if len(fields) == 1 else tuple(res)18 19    else: # or it must be the img directly20        assert fields == ('img',) and isinstance(data, (np.ndarray, Image.Image)), \21            f"data should be an image, not {type(data)}!"22        return data23 24 25def update( data, **fields):26    ''' Called to update the img_and_label27    '''28    if isinstance( data, dict):29        if 'homography' in fields and 'homography' in data:30            data['homography'] = fields.pop('homography') @ data['homography']31        data.update(fields)32        if 'img' in fields: 33            data['imsize'] = data['img'].size34        return data35 36    else: # or it must be the img directly37        return fields['img']38 39 40def rand_log_uniform(rng, a, b):41    return np.exp(rng.uniform(np.log(a),np.log(b)))42 43 44def translate(tx, ty):45    return np.float32(((1,0,tx),(0,1,ty,),(0,0,1)))46 47def rotate(angle):48    return np.float32(((np.cos(angle),-np.sin(angle),0),(np.sin(angle),np.cos(angle),0),(0,0,1)))49 50 51def is_pil_image(img):52    return isinstance(img, Image.Image)53 54 55def homography_from_4pts(pts_cur, pts_new):56    "pts_cur and pts_new = 4x2 point array, in [(x,y),...] format"57    matrix = []58    for p1, p2 in zip(pts_new, pts_cur):59        matrix.append([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])60        matrix.append([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])61    A = np.matrix(matrix, dtype=np.float)62    B = np.array(pts_cur).reshape(8)63 64    homography = np.dot(np.linalg.pinv(A), B)65    homography = tuple(np.array(homography).reshape(8))66    #print(homography)67    return homography68 69 70 71 72