def get_set(folder, iterations, repeats, learning=True): pygame.init() screen = pygame.display.set_mode(DISPLAY) pygame.display.set_caption("Pictures") background = Surface(DISPLAY) background.fill(Color("#000000")) p = Picture() p.load('pictures/samples/car/01.txt') raw_set = generate_pics(['bird', 'cactus', 'car', 'cup', 'helicopter', 'loco', 'octopus', 'tree'], 8) #raw_set = generate_pics(['loco', 'octopus', 'tree'], 8) #getting learning set if learning: p_name = 'train-images.idx3-ubyte' l_name = 'train-labels.idx1-ubyte' else: p_name = 't10k-images.idx3-ubyte' l_name = 't10k-labels.idx1-ubyte' p_f = open(folder+'/'+p_name, 'wb') l_f = open(folder+'/'+l_name, 'wb') l_f.write(struct.pack('>i', 48*4*repeats)) p_f.write(struct.pack('>i', 48*4*repeats)) p_f.write(struct.pack('>i', 32)) p_f.write(struct.pack('>i', 32)) for label in raw_set: for img in raw_set[label]: p = img for r in range(3*repeats): pic = Picture(p.body) for i in range(iterations): pic = add_noise(pic, 40) pic = add_noise(pic, -40) for row in pic.body: for i in row: p_f.write(struct.pack('B', i)) l_f.write(struct.pack('B', label)) pic.show(screen, PIXEL_SIZE) pygame.display.update() pygame.image.save(screen, 'screen'+str(label)+'.png') for label in raw_set: for img in raw_set[label]: p = img for r in range(repeats): pic = Picture(p.body) for i in range(iterations): pic = add_noise(pic, 40) pic = add_noise(pic, -40) for row in pic.body: for i in row: p_f.write(struct.pack('B', i)) l_f.write(struct.pack('B', label)) pic.show(screen, PIXEL_SIZE) pygame.display.update() pygame.image.save(screen, 'screen'+str(r)+'.png') p_f.close() l_f.close()
def generate_pics(folders, number): files = {} for folder in range(len(folders)): files[folder] = [] for n in range(number): files[folder].append('pictures/samples/'+folders[folder]+'/0'+str(n+1)+'.txt') models = {} for a in range(len(folders)): models[a] = [] for file in files[a]: p = Picture() p.load(file) models[a].append(p) p = Picture() p.load(file) p.mirror() models[a].append(p) return models