Spaces:
Sleeping
Sleeping
File size: 1,812 Bytes
9aed787 56ba131 9aed787 56ba131 9aed787 56ba131 9aed787 56ba131 9aed787 4b6da58 9aed787 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
from PIL import Image
from grip_env.environment import GridWorldEnv
import os
import json
import numpy as np
class GenerateBoard():
def __init__(self, level: str, board_size: int, board_number: int):
self.level = level
self.board_size = board_size
metadata_path = os.path.join('src', f'test_{level}.json')
with open(metadata_path, 'r') as f:
metadata = json.load(f)
self.board_num = board_number
self.board_data = metadata[self.board_num]
def setup_initial_board(self):
metadata_obj = self.board_data[-1]
default_start_pos = np.array(metadata_obj['agent_start_pos'])
default_target_pos = np.array(metadata_obj['target_pos'])
info = metadata_obj['info']
target_options = []
for piece in info:
target_info_dict = {}
target = f"{piece['piece_colour']} {piece['piece_shape']} at {piece['piece_region']}"
target_info_dict['target_str'] = target
target_info_dict['piece_info'] = piece
target_options.append(target_info_dict)
env = GridWorldEnv(render_mode="rgb_array", size=self.board_size, grid_info=info, agent_pos=default_start_pos, target_pos=default_target_pos)
env.reset()
image = env.render()
image = Image.fromarray(image)
# Convert all white pixels to gray
image = image.convert("RGBA") # Ensure the image has an alpha channel
data = np.array(image) # Convert image to numpy array
# Create a mask for white pixels
white_pixels = (data[:, :, :3] == [255, 255, 255]).all(axis=2)
data[white_pixels] = [200, 200, 200, 20]
image = Image.fromarray(data) # Convert back to image
return image, target_options, info
|