class TwoPhaseModel:
def __init__(self, MI, NI, granulometry, matrixLabel):
self.MI = MI
self.NI = NI
self.granulometry = granulometry
self.matrixLabel = matrixLabel
xv = np.linspace(0, self.MI - 1, self.MI,
endpoint=True).astype(np.int32)
yv = np.linspace(0, self.NI - 1, self.NI,
endpoint=True).astype(np.int32)
X, Y = np.int32(np.meshgrid(xv, yv))
self.coords = {(x, y) for x, y in zip(X.ravel(), Y.ravel())}
depth = 4
self.qtree = Quadtree(int(depth), Rect(0, 0, int(self.MI),
int(self.NI)))
(self.XX, self.YY) = np.meshgrid(range(0, self.NI), range(0, self.MI))
def compute(self):
Objs = []
Image = np.ones((self.MI, self.NI), np.int32) * self.matrixLabel
Image = Image.astype(np.int32)
start = time.time()
for ix, value in enumerate(
zip(self.granulometry.a_reversed,
self.granulometry.c_reversed)):
print(ix, value[0], value[1])
area_set = 0
while area_set < value[1]:
cy, cx = self.coords.pop()
#if Image[cy,cx] == self.matrixLabel:
if True:
b = self.granulometry.getB(value[0])
theta = random.uniform(0, np.pi)
c = Ellipse(cy, cx, int(value[0]), int(b), theta)
objs = self.qtree.query(c)
if len(objs) == 0:
self.qtree.insert(c)
Objs.append(c)
area_set += c.area()
#ellipseMatrix(c.y(), c.x(), c.a(), c.b(), c.theta(), Image, int(self.granulometry.Label), self.XX, self.YY)
#ellipseDiscard(c.y(), c.x(), c.a(), c.b(), c.theta(), self.XX, self.YY, self.coords, Image, self.granulometry.Label)
else:
self.coords.add((cy, cx))
print(time.time() - start)
return Objs, int(self.granulometry.Label), self.XX, self.YY, Image
def main():
pygame.init()
screen = pygame.display.set_mode((width, height))
map = pygame.Surface((map_width, map_height))
pygame.display.set_caption("quadtree")
font = pygame.font.SysFont('arial', 20)
clock = pygame.time.Clock()
particles = Group([
Particle(
random.gauss(map_width / 2, 150) % map_width,
random.gauss(map_height / 2, 150) % map_height, 4, map)
for _ in range(1000)
])
qtree = Quadtree(map.get_rect(), 4, map)
for p in particles:
qtree.insert(p)
while True:
clock.tick(60)
fps_text = font.render(f'fps:{int(clock.get_fps())}', True,
THECOLORS['white'])
fps_rect = fps_text.get_rect()
for event in pygame.event.get():
if event.type == pygame.QUIT:
exit()
qtree = Quadtree(map.get_rect(), 4, map)
for p in particles:
qtree.insert(p)
particles.update()
for p in particles:
p.highlight = False
query_rect = Rect(p.rect.x, p.rect.y, p.radius * 2, p.radius * 2)
if len(qtree.query(query_rect)) != 1:
p.highlight = True
map.fill(THECOLORS['black'])
particles.draw(map)
pygame.transform.scale(map, (width, height), screen)
screen.blit(fps_text, fps_rect)
pygame.display.update()
def test():
pygame.init()
screen = pygame.display.set_mode((W, H), 0)
quadtree = Quadtree((0, 0, W, H))
triangles = []
p = (0, 0)
while True:
new_triangles = False
new_point = False
update = False
events = pygame.event.get()
for e in events:
if e.type == pygame.QUIT:
return
elif e.type == pygame.KEYUP and e.key == pygame.K_ESCAPE:
return
elif e.type == pygame.KEYUP and e.key == pygame.K_SPACE:
new_triangles = True
elif e.type == pygame.KEYUP and e.key == pygame.K_RETURN:
new_point = True
elif e.type == pygame.MOUSEBUTTONUP:
p = e.pos
update = True
print p
if new_triangles:
for i in range(N_triangles):
quadtree.remove(i)
print "foo"
triangles = []
for i in range(N_triangles):
triangles.append(get_random_triangle())
quadtree.add(i, triangles[-1])
print triangles[-1]
if new_point:
p = get_random_point()
if new_point or new_triangles or update:
screen.fill(BLACK)
ids = quadtree.query(p)
print ids
for i in range(N_triangles):
color = GREEN if i in ids else BLUE
pygame.draw.lines(screen, color, True, triangles[i])
pygame.draw.line(screen, RED, (p[0]-2, p[1]-2), (p[0]+2, p[1]+2))
pygame.draw.line(screen, RED, (p[0]-2, p[1]+2), (p[0]+2, p[1]-2))
pygame.display.update()
