__author__ = 'Miguel Molero' import numpy as np from scipy.misc import imread, imrotate from SimNDT.core.constants import * from SimNDT.core.boundary import Boundary class Scenario: def __init__(self, Width = 40, Height = 40, Pixel_mm =10, Label=0): self.Width = Width self.Height = Height self.Pixel_mm = Pixel_mm self.Label = Label self.M = int(self.Height * self.Pixel_mm) self.N = int(self.Width * self.Pixel_mm) self.I = np.ones((self.M,self.N), dtype=np.uint8)*Label self.Iabs = 0 self.Io = np.ones((self.M,self.N), dtype=np.uint8)*Label self.Tap = list() self.BC = False def setImage(self,I, Width, Height, Pixel_mm, Label): self.I = np.copy(I).astype(np.uint8) self.Io = np.copy(I).astype(np.uint8) self.Width = Width self.Height = Height self.Pixel_mm = Pixel_mm self.Label = Label self.M = int(self.Height * self.Pixel_mm) self.N = int(self.Width * self.Pixel_mm) self.resetBoundary() def resetBoundary(self): self.Iabs = 0 self.Tap = list() self.BC = False def __str__(self): return "Scenario: " def __repr__(self): return "Scenario: " def createBoundaries(self,boundaries): self.M, self.N = np.shape(self.I) self.Width = int(self.N/float(self.Pixel_mm)) self.Height = int(self.N/float(self.Pixel_mm)) self.resetBoundary() for boundary in boundaries: if boundary.BC == BC.AirLayer: size = 1 else: size = boundary.Size * self.Pixel_mm if boundary.Name == "Top": topSize = size elif boundary.Name == "Bottom": bottomSize = size elif boundary.Name == "Left": leftSize = size elif boundary.Name == "Right": rightSize = size self.Tap = np.array([topSize, bottomSize, leftSize, rightSize]) self.Iabs = self.applyBoundaries(self.I) self.BC = True def applyBoundaries(self, I): self.Tap = np.int32(self.Tap) M_abs = int( self.M + self.Tap[0] + self.Tap[1] ) N_abs = int( self.N + self.Tap[2] + self.Tap[3] ) Iabs = 255*np.ones((int(M_abs),int(N_abs)),dtype=np.uint8) Iabs[self.Tap[0] : M_abs-self.Tap[1], self.Tap[2] : N_abs-self.Tap[3]] = np.copy(I) return Iabs def updateScenario(self): if self.BC: self.Iabs = self.applyBoundaries(self.I) def addEllipse(self, x0, y0, a, b, theta, Label): x0 *= self.Pixel_mm y0 *= self.Pixel_mm a *= self.Pixel_mm b *= self.Pixel_mm (x,y) = np.meshgrid(range(0,self.N),range(0,self.M)) Ellipse = ( ( ( (x-x0)*np.cos(theta)+(y-y0)*np.sin(theta) )**2 )/(a**2) + ( ( (x-x0)*np.sin(theta)-(y-y0)*np.cos(theta) )**2 )/(b**2) ) Img = (Ellipse < 1.0) indx,indy = np.nonzero(Img == 1) self.I[indx,indy] = Label self.updateScenario() def resetImage(self): self.resetBoundary() self.I = np.copy(self.Io) def addObject(self, obj): if obj.Name == "ellipse": a = obj.a b = obj.b theta = obj.theta Label = obj.Label x0 = obj.x0 * self.Pixel_mm y0 = obj.y0 * self.Pixel_mm self.addEllipse(x0, y0, a, b, theta, Label) elif obj.Name == "circle": x0 = obj.x0 y0 = obj.y0 a = obj.r Label = obj.Label self.addEllipse(x0, y0, a, a, 0, Label) elif obj.Name == "square": x0 = obj.x0 y0 = obj.y0 L = obj.L theta = obj.theta Label = obj.Label self.addRectangle(x0, y0, L, L, theta, Label) elif obj.Name == "rectangle": x0 = obj.x0 y0 = obj.y0 W = obj.W H = obj.H theta = obj.theta Label = obj.Label self.addRectangle(x0, y0, W, H, theta, Label) # @blab+ # Make a copy of I, just for visualization (I is graphics view and material matrix for EDIF) def addTransducer(self, x0, y0, Size, Label): ox = 0 oy = 0 if len(self.Tap) != 0: ox = self.Tap[2] oy = self.Tap[0] x0 = x0+ox y0 = y0+oy a = Size b = a theta = 0 x0 *= self.Pixel_mm y0 *= self.Pixel_mm a *= self.Pixel_mm b *= self.Pixel_mm (x,y) = np.meshgrid(range(0,self.N),range(0,self.M)) Ellipse = ( ( ( (x-x0)*np.cos(theta)+(y-y0)*np.sin(theta) )**2 )/(a**2) + ( ( (x-x0)*np.sin(theta)-(y-y0)*np.cos(theta) )**2 )/(b**2) ) Img = (Ellipse < 1.0) if np.size(self.Iabs)==1: I = np.copy(self.I) else: I = np.copy(self.Iabs) indx,indy = np.nonzero(Img == 1) I[indx,indy] = Label return I def addSensor(self, I, x0, y0, Size, Label): ox = 0 oy = 0 if len(self.Tap) != 0: ox = self.Tap[2] oy = self.Tap[0] x0 = x0+ox y0 = y0+oy a = Size b = a theta = 0 x0 *= self.Pixel_mm y0 *= self.Pixel_mm a *= self.Pixel_mm b *= self.Pixel_mm (x,y) = np.meshgrid(range(0,self.N),range(0,self.M)) Ellipse = ( ( ( (x-x0)*np.cos(theta)+(y-y0)*np.sin(theta) )**2 )/(a**2) + ( ( (x-x0)*np.sin(theta)-(y-y0)*np.cos(theta) )**2 )/(b**2) ) Img = (Ellipse < 1.0) indx,indy = np.nonzero(Img == 1) I[indx,indy] = Label return I def addRectangle(self, W_0, H_0, W, H, Theta,Label): """ Create a rectangle W_0, H_0: center of rectangle W,H : dimension of rectangle """ a = round(H*self.Pixel_mm/2.0) b = round(W*self.Pixel_mm/2.0) angle = Theta * np.pi/180.0 H0 = H_0*self.Pixel_mm W0 = W_0*self.Pixel_mm if angle == 0: size_a = 4*a size_b = 4*b vectX = np.int32(np.linspace(-a,a, size_a) + H0) vectY = np.int32(np.linspace(-b,b, size_b) + W0) for x in vectX: for y in vectY: self.I[x, y] = Label else: size_a = 4*a size_b = 4*b for x in np.linspace(-a,a,size_a): for y in np.linspace(-b,b,size_b): _xr = np.int32(np.cos(angle)*x - np.sin(angle)*y + H0) _yr = np.int32(np.sin(angle)*x + np.cos(angle)*y + W0) self.I[_xr,_yr] = Label self.updateScenario() def rotate(self, angle=90, direction="clockwise"): if direction == "clockwise": self.I = imrotate(self.I,-1*angle, interp = 'nearest') self.M, self.N = np.shape(self.I) if np.size(self.Iabs) > 1: self.Iabs = imrotate(self.Iabs,-1*angle, interp = 'nearest') else: self.I = imrotate(self.I,angle, interp = 'nearest') self.M, self.N = np.shape(self.I) if np.size(self.Iabs) > 1: self.Iabs = imrotate(self.Iabs,angle, interp = 'nearest')