Volovikov_CV/falling_ball/main.py

import pygame
import pymunk.pygame_util
from cv import recognize
import cv2

# image = cv2.imread('falling_ball/image.png')
cam = cv2.VideoCapture(0)

cap, image = cam.read()

RES = WIDTH, HEIGHT = 1280, 800
BALL_RADIUS = min(RES) / 30
FPS = 60

polys = recognize(image, RES)

pygame.init()

surface = pygame.display.set_mode((0, 0), pygame.FULLSCREEN, display=1)
clock = pygame.time.Clock()

draw_options = pymunk.pygame_util.DrawOptions(surface)
space = pymunk.Space()
space.gravity = 0, 2000

def create_ball(space, pos: tuple[int ,int]):
    ball_mass, ball_radius = 10, BALL_RADIUS
    ball_moment = pymunk.moment_for_circle(ball_mass, 0, ball_radius)
    ball_body = pymunk.Body(ball_mass, ball_moment)
    ball_body.position = pos
    ball_shape = pymunk.Circle(ball_body, ball_radius)
    ball_shape.elasticity = 0.2
    space.add(ball_body, ball_shape)

def draw_by_lines(space, lines=list[tuple[int, int]]):
    for line_i in range(len(lines)):
        line = pymunk.Segment(space.static_body, lines[line_i], lines[(line_i+1)%4], 2)
        line.elasticity = 0.1
        space.add(line)

for poly in polys:
    draw_by_lines(space, poly)

circle_position = (350, BALL_RADIUS+1)

while True:
    surface.fill(pygame.Color('white'))

    pygame.draw.circle(surface, pygame.Color('BLUE'), circle_position, BALL_RADIUS)

    for i in pygame.event.get():
        if i.type == pygame.KEYDOWN:
            if i.key == pygame.K_SPACE:
                create_ball(space, circle_position)
            if i.key == pygame.K_RIGHT:
                circle_position = (circle_position[0] + 10, circle_position[1])
            if i.key == pygame.K_LEFT:
                circle_position = (circle_position[0] - 10, circle_position[1])
        if i.type == pygame.QUIT:
            exit()

    space.step(1 / FPS)
    space.debug_draw(draw_options)
    pygame.display.flip()
    clock.tick(FPS)