#coding:utf-8
import random
import cv2
import math
import numpy as np

class DataExtention:

    def dirt(self,image):
        _image = image
        _size = image.shape[:2] #height:width
        _c = random.randint(0,10)
        _count = 0
        _r = 1
        while _count <= _c :
            _x = random.randint(0,_size[1]-1)
            _y = random.randint(0,_size[0]-1)
            _image = cv2.rectangle(_image,(_x,_y),(_x+_r,_y+_r),(255,255,255),-1)
            _count += 1
        return _image

    def rotate(self,image):
        _image = image
        _angle = random.randint(-10,10)
        _rad = math.radians(_angle)
        _size = (_image.shape[1],_image.shape[0])
        if _angle<0:
            upset = (0,0)
        else:
            upset = (_size[0],_size[1])
        _x = (int)(math.fabs(_size[0]*math.cos(_rad)))
        _y = (int)(math.fabs(_size[0]*math.sin(_rad))+math.fabs(_size[1]*math.cos(_rad)))
        scale = 1.0
        rotation_matrix = cv2.getRotationMatrix2D(upset,_angle,scale)
        _image = cv2.warpAffine(_image,rotation_matrix,(_x,_y))
        
        return _image

    def transform(self,image):
        _x = random.randint(-5,5)
        _y = random.randint(-5,5)
        _image = image
        _size = (_image.shape[1],_image.shape[0])
        _transM = np.float32([[1,0,_x],[0,1,_y]])
        _image = cv2.warpAffine(_image,_transM,(_size[0],_size[1]))

        return _image