Video-Gen / modules /particles.py
jacky3102's picture
Create modules/particles.py
915a172 verified
Raw
History Blame Contribute Delete
3.67 kB
import random
import math
import cv2
import numpy as np
from PIL import Image, ImageDraw
class Particle:
def __init__(self, canvas_w, canvas_h, effect_type):
self.canvas_w, self.canvas_h = canvas_w, canvas_h
self.effect_type = effect_type
self.reset()
def reset(self):
self.x = random.uniform(0, self.canvas_w)
self.y = random.uniform(-50, self.canvas_h + 50)
self.size = random.uniform(2, 6)
self.speed = random.uniform(0.5, 2.5)
self.wind = random.uniform(-0.5, 0.5)
self.opacity = random.uniform(100, 255)
self.rotation = random.uniform(0, 360)
self.rot_speed = random.uniform(-3, 3)
self.wobble = random.uniform(0, math.pi * 2)
def update(self):
self.y += self.speed
self.x += self.wind + math.sin(self.wobble) * 0.3
self.wobble += 0.02
self.rotation += self.rot_speed
if self.y > self.canvas_h + 50:
self.reset()
if self.x < -50:
self.x = self.canvas_w + 50
if self.x > self.canvas_w + 50:
self.x = -50
def draw_snowflake(draw, x, y, size, opacity):
r = int(size * 3)
color = (255, 255, 255, int(opacity))
draw.ellipse([x - r, y - r, x + r, y + r], fill=color)
for angle in [0, math.pi / 3, math.pi * 2 / 3]:
ex, ey = x + math.cos(angle) * r * 1.5, y + math.sin(angle) * r * 1.5
draw.line([x, y, ex, ey], fill=color, width=max(1, int(size)))
def draw_leaf(draw, x, y, size, opacity, rotation):
leaf_colors = [(180, 60, 30), (200, 120, 20), (220, 160, 40), (160, 80, 20)]
color = random.choice(leaf_colors) + (int(opacity),)
r = int(size * 5)
img = Image.new('RGBA', (r * 2, r * 2), (0, 0, 0, 0))
leaf_draw = ImageDraw.Draw(img)
leaf_draw.ellipse([r * 0.2, r * 0.5, r * 1.8, r * 1.5], fill=color)
img = img.rotate(rotation, expand=False)
draw._image.paste(img, (int(x - r), int(y - r)), img)
def draw_raindrop(draw, x, y, size, opacity):
color = (180, 200, 240, int(opacity))
r = int(size * 1.5)
draw.ellipse([x - r, y - r * 3, x + r, y + r * 3], fill=color)
def draw_sparkle(draw, x, y, size, opacity):
gold = (255, 215, 0, int(opacity))
r = int(size * 2)
draw.ellipse([x - r, y - r, x + r, y + r], fill=gold)
for angle in range(0, 360, 45):
rad = math.radians(angle)
ex, ey = x + math.cos(rad) * r * 2, y + math.sin(rad) * r * 2
draw.line([x, y, ex, ey], fill=gold, width=max(1, int(size * 0.8)))
def draw_dust(draw, x, y, size, opacity):
color = (200, 190, 170, int(opacity))
r = int(size * 1.2)
draw.ellipse([x - r, y - r, x + r, y + r], fill=color)
def apply_effect(frame, particles, effect_type):
img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGBA))
overlay = Image.new('RGBA', img.size, (0, 0, 0, 0))
draw = ImageDraw.Draw(overlay)
draw._image = overlay
for p in particles:
p.update()
opacity = max(10, min(255, p.opacity * (1 - (p.y / (p.canvas_h + 50)))))
if effect_type == "snow": draw_snowflake(draw, p.x, p.y, p.size, opacity)
elif effect_type == "leaves": draw_leaf(draw, p.x, p.y, p.size, opacity, p.rotation)
elif effect_type == "rain": draw_raindrop(draw, p.x, p.y, p.size, opacity)
elif effect_type == "sparkle": draw_sparkle(draw, p.x, p.y, p.size, opacity)
elif effect_type == "dust": draw_dust(draw, p.x, p.y, p.size, opacity)
img = Image.alpha_composite(img, overlay)
return cv2.cvtColor(np.array(img), cv2.COLOR_RGBA2BGR)
EFFECTS = ["snow", "leaves", "rain", "sparkle", "dust", "flame", "heat", "oldfilm", "none"]