SmokeSim/examples/image_augmentation.py at master · q-viper/SmokeSim · GitHub

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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
from smokesim.augmentation import Augmentation
from smokesim.defs import SmokeProperty, ParticleProperty

import numpy as np
from pathlib import Path

if __name__ == "__main__":
    WIDTH, HEIGHT = 700, 500
    project_dir = Path(__file__).resolve().parents[1]
    img_path = project_dir / Path("assets/me.jpg")
    print(f"Image path: {img_path}")
    augmentation = Augmentation(image_path=img_path, screen_dim=(WIDTH, HEIGHT))
    smoke_machine = augmentation.smoke_machine
    smoke_properties = SmokeProperty(
        particle_count=15,
        sprite_size=25,
        origin=(250, 500),
    )
    augmentation.add_smoke(smoke_properties)
    smoke_properties = SmokeProperty(
        particle_count=15,
        sprite_size=25,
        origin=(450, 500),
    )
    augmentation.add_smoke(smoke_properties)
    random_state = np.random.RandomState(42)

    final_image = augmentation.augment(
        steps=90, history_path=project_dir / Path("media/smoke_history.mp4")
    )
    final_image, final_mask = augmentation.augment(steps=90)
    for i in range(5):
        smoke_properties = SmokeProperty(
            color=smoke_machine.color,
            particle_count=1,
            origin=(
                random_state.randint(100, WIDTH),
                random_state.randint(100, HEIGHT),
            ),
            lifetime=200,
            particle_property=ParticleProperty(
                min_lifetime=200,
                max_lifetime=500,
                min_scale=10,
                max_scale=50,
                fade_speed=50,
                scale=50,
                smoke_sprite_size=50,
                color=smoke_machine.color,
            ),
        )
        augmentation.add_smoke(
            smoke_property=smoke_properties,
        )
    final_image1, final_mask1 = augmentation.augment(steps=1)
    augmentation.save_as(project_dir / Path("assets/augmented_smoke_image.png"))
    augmentation.end()

    augmentation = Augmentation(image_path=img_path, screen_dim=(WIDTH, HEIGHT))
    smoke_machine = augmentation.smoke_machine
    smoke_properties = SmokeProperty(
        particle_count=15,
        sprite_size=25,
        origin=(250, 500),
    )
    augmentation.add_smoke(smoke_properties)
    smoke_properties = SmokeProperty(
        particle_count=15,
        sprite_size=25,
        origin=(450, 500),
    )
    augmentation.add_smoke(smoke_properties)
    random_state = np.random.RandomState(42)

    final_image = augmentation.augment(
        steps=90, history_path=project_dir / Path("media/smoke_history.mp4")
    )
    final_image, final_mask = augmentation.augment(steps=90)
    for i in range(5):
        smoke_properties = SmokeProperty(
            color=smoke_machine.color,
            particle_count=1,
            origin=(
                random_state.randint(100, WIDTH),
                random_state.randint(100, HEIGHT),
            ),
            lifetime=200,
            particle_property=ParticleProperty(
                min_lifetime=200,
                max_lifetime=500,
                min_scale=10,
                max_scale=50,
                fade_speed=50,
                scale=50,
                smoke_sprite_size=50,
                color=smoke_machine.color,
            ),
        )
        augmentation.add_smoke(
            smoke_property=smoke_properties,
        )
    final_image2, final_mask2 = augmentation.augment(steps=1)
    augmentation.save_as(project_dir / Path("assets/augmented_smoke_image.png"))
    augmentation.end()

    print(f"Are equal: {np.array_equal(final_image1, final_image2)}")