importioimportosimporttime#import tkinter as tkimportnumpyasnpfromPILimportImage,ImageDraw,ImageFontfromnumpyimportarrayfrompkg_resourcesimportresource_stringasresource_bytesfromflatland.utils.graphics_layerimportGraphicsLayerfromflatland.core.grid.rail_env_gridimportRailEnvTransitions# noqa: E402
[docs]classPILGL(GraphicsLayer):# tk.Tk() must be a singleton!# https://stackoverflow.com/questions/26097811/image-pyimage2-doesnt-exist# window = tk.Tk()RAIL_LAYER=0PREDICTION_PATH_LAYER=1TARGET_LAYER=2AGENT_LAYER=3SELECTED_AGENT_LAYER=4SELECTED_TARGET_LAYER=5def__init__(self,width,height,jupyter=False,screen_width=800,screen_height=600):self.yxBase=(0,0)self.linewidth=4self.n_agent_colors=1# overridden in loadAgentself.width=widthself.height=heightself.background_grid=np.zeros(shape=(self.width,self.height))ifjupyterisFalse:# NOTE: Currently removed the dependency on# screeninfo. We have to find an alternate# way to compute the screen width and height# In the meantime, we are harcoding the 800x600# assumptionself.screen_width=screen_widthself.screen_height=screen_heightw=(self.screen_width-self.width-10)/(self.width+1+self.linewidth)h=(self.screen_height-self.height-10)/(self.height+1+self.linewidth)self.nPixCell=int(max(1,np.ceil(min(w,h))))else:self.nPixCell=40# Total grid size at native scaleself.widthPx=self.width*self.nPixCell+self.linewidthself.heightPx=self.height*self.nPixCell+self.linewidthself.xPx=int((self.screen_width-self.widthPx)/2.0)self.yPx=int((self.screen_height-self.heightPx)/2.0)self.layers=[]self.draws=[]self.tColBg=(255,255,255)# white backgroundself.tColRail=(0,0,0)# black railsself.tColGrid=(230,)*3# light grey for gridsColors="d50000#c51162#aa00ff#6200ea#304ffe#2962ff#0091ea#00b8d4#00bfa5#00c853"+ \
"#64dd17#aeea00#ffd600#ffab00#ff6d00#ff3d00#5d4037#455a64"self.agent_colors=[self.rgb_s2i(sColor)forsColorinsColors.split("#")]self.n_agent_colors=len(self.agent_colors)self.firstFrame=Trueself.old_background_image=(None,None,None)self.create_layers()self.font=ImageFont.load_default()
[docs]defbuild_background_map(self,dTargets):x=self.old_background_imagerebuild=Falseifx[0]isNone:rebuild=Trueelse:iflen(x[0])!=len(dTargets):rebuild=Trueelse:ifx[0]!=dTargets:rebuild=Trueifx[1]!=self.width:rebuild=Trueifx[2]!=self.height:rebuild=Trueifrebuild:# rebuild background_grid to control the visualisation of buildings, trees, mountains, lakes and riverself.background_grid=np.zeros(shape=(self.width,self.height))# build base distance map (distance to targets)forxinrange(self.width):foryinrange(self.height):distance=int(np.ceil(np.sqrt(self.width**2.0+self.height**2.0)))forrcindTargets:r=rc[1]c=rc[0]d=int(np.floor(np.sqrt((x-r)**2+(y-c)**2))/0.5)distance=min(d,distance)self.background_grid[x][y]=distanceself.old_background_image=(dTargets,self.width,self.height)
[docs]defrgb_s2i(self,sRGB):""" convert a hex RGB string like 0091ea to 3-tuple of ints """returntuple(int(sRGB[iRGB*2:iRGB*2+2],16)foriRGBin[0,1,2])
[docs]defplot(self,gX,gY,color=None,linewidth=3,layer=RAIL_LAYER,opacity=255,**kwargs):""" Draw a line joining the points in gX, GY - each an"""color=self.adapt_color(color)iflen(color)==3:color+=(opacity,)eliflen(color)==4:color=color[:3]+(opacity,)gPoints=np.stack([array(gX),-array(gY)]).T*self.nPixCellgPoints=list(gPoints.ravel())# the width here was self.linewidth - not really sure of the implicationsself.draws[layer].line(gPoints,fill=color,width=linewidth)
[docs]defdraw_image_xy(self,pil_img,xyPixLeftTop,layer=RAIL_LAYER,):# Resize all PIL images just before drawing them# to ensure that resizing doesnt affect the# recolorizing strategies in place## That said : All the code in this file needs# some serious refactoring -_- to ensure the# code style and structure is consitent.# - Mohantypil_img=pil_img.resize((self.nPixCell,self.nPixCell))if(pil_img.mode=="RGBA"):pil_mask=pil_imgelse:pil_mask=Noneself.layers[layer].paste(pil_img,xyPixLeftTop,pil_mask)
[docs]defbegin_frame(self):# Create a new agent layerself.create_layer(iLayer=PILGL.AGENT_LAYER,clear=True)self.create_layer(iLayer=PILGL.PREDICTION_PATH_LAYER,clear=True)
[docs]defget_image(self):""" return a blended / alpha composited image composed of all the layers, with layer 0 at the "back". """img=self.alpha_composite_layers()returnarray(img)
[docs]defsave_image(self,filename):""" Renders the current scene into a image file :param filename: filename where to store the rendering output_generator (supported image format *.bmp , .. , *.png) """img=self.alpha_composite_layers()img.save(filename)
[docs]defclear_layer(self,iLayer=0,opacity=None):ifopacityisNone:opacity=0ifiLayer>0else255self.layers[iLayer]=img=self.create_image(opacity)# We also need to maintain a Draw object for each layerself.draws[iLayer]=ImageDraw.Draw(img)
[docs]defcreate_layer(self,iLayer=0,clear=True):# If we don't have the layers already, create themiflen(self.layers)<=iLayer:foriinrange(len(self.layers),iLayer+1):ifi==0:opacity=255# "bottom" layer is opaque (for rails)else:opacity=0# subsequent layers are transparentimg=self.create_image(opacity)self.layers.append(img)self.draws.append(ImageDraw.Draw(img))else:# We do already have this iLayer. Clear it if requested.ifclear:self.clear_layer(iLayer)
[docs]defcreate_layers(self,clear=True):self.create_layer(PILGL.RAIL_LAYER,clear=clear)# rail / background (scene)self.create_layer(PILGL.AGENT_LAYER,clear=clear)# agentsself.create_layer(PILGL.TARGET_LAYER,clear=clear)# agentsself.create_layer(PILGL.PREDICTION_PATH_LAYER,clear=clear)# drawing layer for agent's prediction pathself.create_layer(PILGL.SELECTED_AGENT_LAYER,clear=clear)# drawing layer for selected agentself.create_layer(PILGL.SELECTED_TARGET_LAYER,clear=clear)# drawing layer for selected agent's target
[docs]classPILSVG(PILGL):""" Note : This class should now ideally be called as PILPNG, but for backward compatibility, and to not introduce any breaking changes at this point we are sticking to the legacy name of PILSVG (when in practice we are not using SVG anymore) """def__init__(self,width,height,jupyter=False,screen_width=800,screen_height=600):oSuper=super()oSuper.__init__(width,height,jupyter,screen_width,screen_height)self.lwAgents=[]self.agents_prev=[]self.load_buildings()self.load_scenery()self.load_rail()self.load_agent()
[docs]defload_rail(self):""" Load the rail SVG images, apply rotations, and store as PIL images. """rail_files={"":"Background_Light_green.png","WE":"Gleis_Deadend.png","WW EE NN SS":"Gleis_Diamond_Crossing.png","WW EE":"Gleis_horizontal.png","EN SW":"Gleis_Kurve_oben_links.png","WN SE":"Gleis_Kurve_oben_rechts.png","ES NW":"Gleis_Kurve_unten_links.png","NE WS":"Gleis_Kurve_unten_rechts.png","NN SS":"Gleis_vertikal.png","NN SS EE WW ES NW SE WN":"Weiche_Double_Slip.png","EE WW EN SW":"Weiche_horizontal_oben_links.png","EE WW SE WN":"Weiche_horizontal_oben_rechts.png","EE WW ES NW":"Weiche_horizontal_unten_links.png","EE WW NE WS":"Weiche_horizontal_unten_rechts.png","NN SS EE WW NW ES":"Weiche_Single_Slip.png","NE NW ES WS":"Weiche_Symetrical.png","NN SS EN SW":"Weiche_vertikal_oben_links.png","NN SS SE WN":"Weiche_vertikal_oben_rechts.png","NN SS NW ES":"Weiche_vertikal_unten_links.png","NN SS NE WS":"Weiche_vertikal_unten_rechts.png","NE NW ES WS SS NN":"Weiche_Symetrical_gerade.png","NE EN SW WS":"Gleis_Kurve_oben_links_unten_rechts.png"}target_files={"EW":"Bahnhof_#d50000_Deadend_links.png","NS":"Bahnhof_#d50000_Deadend_oben.png","WE":"Bahnhof_#d50000_Deadend_rechts.png","SN":"Bahnhof_#d50000_Deadend_unten.png","EE WW":"Bahnhof_#d50000_Gleis_horizontal.png","NN SS":"Bahnhof_#d50000_Gleis_vertikal.png"}# Dict of rail cell images indexed by binary transitionspil_rail_files_org=self.load_pngs(rail_files,rotate=True)pil_rail_files=self.load_pngs(rail_files,rotate=True,background_image="Background_rail.png",whitefilter="Background_white_filter.png")# Load the target files (which have rails and transitions of their own)# They are indexed by (binTrans, iAgent), ie a tuple of the binary transition and the agent indexpil_target_files_org=self.load_pngs(target_files,rotate=False,agent_colors=self.agent_colors)pil_target_files=self.load_pngs(target_files,rotate=False,agent_colors=self.agent_colors,background_image="Background_rail.png",whitefilter="Background_white_filter.png")# Load station and recolorize themstation=self.pil_from_png_file('flatland.png',"Bahnhof_#d50000_target.png")self.station_colors=self.recolor_image(station,[0,0,0],self.agent_colors,False)cell_occupied=self.pil_from_png_file('flatland.png',"Cell_occupied.png")self.cell_occupied=self.recolor_image(cell_occupied,[0,0,0],self.agent_colors,False)# Merge them with the regular rails.# https://stackoverflow.com/questions/38987/how-to-merge-two-dictionaries-in-a-single-expressionself.pil_rail={**pil_rail_files,**pil_target_files}self.pil_rail_org={**pil_rail_files_org,**pil_target_files_org}
[docs]defload_pngs(self,file_directory,rotate=False,agent_colors=False,background_image=None,whitefilter=None):pil={}transitions=RailEnvTransitions()directions=list("NESW")fortransition,fileinfile_directory.items():# Translate the ascii transition description in the format "NE WS" to the# binary list of transitions as per RailEnv - NESW (in) x NESW (out)transition_16_bit=["0"]*16forsTranintransition.split(" "):iflen(sTran)==2:in_direction=directions.index(sTran[0])out_direction=directions.index(sTran[1])transition_idx=4*in_direction+out_directiontransition_16_bit[transition_idx]="1"transition_16_bit_string="".join(transition_16_bit)binary_trans=int(transition_16_bit_string,2)pil_rail=self.pil_from_png_file('flatland.png',file).convert("RGBA")ifbackground_imageisnotNone:img_bg=self.pil_from_png_file('flatland.png',background_image).convert("RGBA")pil_rail=Image.alpha_composite(img_bg,pil_rail)ifwhitefilterisnotNone:img_bg=self.pil_from_png_file('flatland.png',whitefilter).convert("RGBA")pil_rail=Image.alpha_composite(pil_rail,img_bg)ifrotate:# For rotations, we also store the base imagepil[binary_trans]=pil_rail# Rotate both the transition binary and the image and save in the dictfornRotin[90,180,270]:binary_trans_2=transitions.rotate_transition(binary_trans,nRot)# PIL rotates anticlockwise for positive thetapil_rail_2=pil_rail.rotate(-nRot)pil[binary_trans_2]=pil_rail_2ifagent_colors:# For recoloring, we don't store the base image.base_color=self.rgb_s2i("d50000")pils=self.recolor_image(pil_rail,base_color,self.agent_colors)forcolor_idx,pil_rail_2inenumerate(pils):pil[(binary_trans,color_idx)]=pils[color_idx]returnpil
[docs]defset_rail_at(self,row,col,binary_trans,target=None,is_selected=False,rail_grid=None,num_agents=None,show_debug=True):ifbinary_transinself.pil_rail:pil_track=self.pil_rail[binary_trans]iftargetisnotNone:target_img=self.station_colors[target%len(self.station_colors)]target_img=Image.alpha_composite(pil_track,target_img)self.draw_image_row_col(target_img,(row,col),layer=PILGL.TARGET_LAYER)ifshow_debug:self.text_rowcol((row+0.8,col+0.0),strText=str(target),layer=PILGL.TARGET_LAYER)city_size=1ifnum_agentsisnotNone:city_size=max(1,np.log(1+num_agents)/2.5)ifbinary_trans==0:ifself.background_grid[col][row]<=4+np.ceil(((col*row+col)%10)/city_size):a=int(self.background_grid[col][row])a=a%len(self.lBuildings)if(col+row+col*row)%13>11:pil_track=self.scenery[a%len(self.scenery)]else:if(col+row+col*row)%3==0:a=(a+(col+row+col*row))%len(self.lBuildings)pil_track=self.lBuildings[a]elif((self.background_grid[col][row]>5+((col*row+col)%3))or((col**3+row**2+col*row)%10==0)):a=int(self.background_grid[col][row])-4a2=(a+(col+row+col*row+col**3+row**4))ifa2%64>11:a=a2a_l=a%len(self.scenery)ifa2%50==49:pil_track=self.scenery_water[0]else:pil_track=self.scenery[a_l]ifrail_gridisnotNone:ifa2%11>3:ifa_l==len(self.scenery)-1:# mountainifcol>1androw%7==1:ifrail_grid[row,col-1]==0:self.draw_image_row_col(self.scenery_d2[0],(row,col-1),layer=PILGL.RAIL_LAYER)pil_track=self.scenery_d2[1]else:ifa_l==len(self.scenery)-1:# mountainifcol>2andnot(row%7==1):ifrail_grid[row,col-2]==0andrail_grid[row,col-1]==0:self.draw_image_row_col(self.scenery_d3[0],(row,col-2),layer=PILGL.RAIL_LAYER)self.draw_image_row_col(self.scenery_d3[1],(row,col-1),layer=PILGL.RAIL_LAYER)pil_track=self.scenery_d3[2]self.draw_image_row_col(pil_track,(row,col),layer=PILGL.RAIL_LAYER)else:print("Can't render - illegal rail or SVG element is undefined:",row,col,format(binary_trans,"#018b")[2:],binary_trans)iftargetisnotNone:ifis_selected:svgBG=self.pil_from_png_file('flatland.png',"Selected_Target.png")self.clear_layer(PILGL.SELECTED_TARGET_LAYER,0)self.draw_image_row_col(svgBG,(row,col),layer=PILGL.SELECTED_TARGET_LAYER)
[docs]defrecolor_image(self,pil,a3BaseColor,ltColors,invert=False):rgbaImg=array(pil)pils=[]foriColor,tnColorinenumerate(ltColors):# find the pixels which match the base paint colorifinvert:xy_color_mask=np.all(rgbaImg[:,:,0:3]-a3BaseColor!=0,axis=2)else:xy_color_mask=np.all(rgbaImg[:,:,0:3]-a3BaseColor==0,axis=2)rgbaImg2=np.copy(rgbaImg)# Repaint the base color with the new colorrgbaImg2[xy_color_mask,0:3]=tnColorpil2=Image.fromarray(rgbaImg2)pils.append(pil2)returnpils
[docs]defload_agent(self):# Seed initial train/zug files indexed by tuple(iDirIn, iDirOut):file_directory={(0,0):"Zug_Gleis_#0091ea.png",(1,2):"Zug_1_Weiche_#0091ea.png",(0,3):"Zug_2_Weiche_#0091ea.png"}# "paint" color of the train images we load - this is the color we will change.# base_color = self.rgb_s2i("0091ea") \# noqa: E800# temporary workaround for trains / agents renamed with different colour:base_color=self.rgb_s2i("d50000")self.pil_zug={}fordirections,path_svginfile_directory.items():in_direction,out_direction=directionspil_zug=self.pil_from_png_file('flatland.png',path_svg)# Rotate both the directions and the image and save in the dictforrot_directioninrange(4):rotation_degree=rot_direction*90in_direction_2=(in_direction+rot_direction)%4out_direction_2=(out_direction+rot_direction)%4# PIL rotates anticlockwise for positive thetapil_zug_2=pil_zug.rotate(-rotation_degree)# Save colored versions of each rotation / variantpils=self.recolor_image(pil_zug_2,base_color,self.agent_colors)forcolor_idx,pil_zug_3inenumerate(pils):self.pil_zug[(in_direction_2,out_direction_2,color_idx)]=pils[color_idx]
[docs]defset_agent_at(self,agent_idx,row,col,in_direction,out_direction,is_selected,rail_grid=None,show_debug=False,clear_debug_text=True,malfunction=False):delta_dir=(out_direction-in_direction)%4color_idx=agent_idx%self.n_agent_colors# when flipping direction at a dead end, use the "out_direction" direction.ifdelta_dir==2:in_direction=out_directionpil_zug=self.pil_zug[(in_direction%4,out_direction%4,color_idx)]self.draw_image_row_col(pil_zug,(row,col),layer=PILGL.AGENT_LAYER)ifrail_gridisnotNone:ifrail_grid[row,col]==0.0:self.draw_image_row_col(self.scenery_background_white,(row,col),layer=PILGL.RAIL_LAYER)ifis_selected:bg_svg=self.pil_from_png_file('flatland.png',"Selected_Agent.png")self.clear_layer(PILGL.SELECTED_AGENT_LAYER,0)self.draw_image_row_col(bg_svg,(row,col),layer=PILGL.SELECTED_AGENT_LAYER)ifshow_debug:ifnotclear_debug_text:dr=0.2dc=0.2ifin_direction==0:dr=0.8dc=0.0ifin_direction==1:dr=0.0dc=0.8ifin_direction==2:dr=0.4dc=0.8ifin_direction==3:dr=0.8dc=0.4self.text_rowcol((row+dr,col+dc,),str(agent_idx),layer=PILGL.SELECTED_AGENT_LAYER)else:self.text_rowcol((row+0.2,col+0.2,),str(agent_idx))ifmalfunction:self.draw_malfunction(agent_idx,(row,col))