lerobot/examples/hopejr/rolling/test.py

75 lines
2.2 KiB
Python

import numpy as np
from PIL import Image, ImageSequence
def coalesce_gif(im):
"""
Attempt to coalesce frames so each one is a full image.
This handles many (though not all) partial-frame GIFs.
"""
# Convert mode to RGBA
im = im.convert("RGBA")
# Prepare an accumulator the same size as the base frame
base = Image.new("RGBA", im.size)
frames = []
# Go through each frame
for frame in ImageSequence.Iterator(im):
base.alpha_composite(frame.convert("RGBA"))
frames.append(base.copy())
return frames
def remove_white_make_black(arr, threshold=250):
"""
For each pixel in arr (H,W,3), if R,G,B >= threshold, set to black (0,0,0).
This effectively 'removes' white so it won't affect the sum.
"""
mask = (arr[..., 0] >= threshold) & \
(arr[..., 1] >= threshold) & \
(arr[..., 2] >= threshold)
arr[mask] = 0 # set to black
def main():
# Load the animated GIF
gif = Image.open("input.gif")
# Coalesce frames so each is full-size
frames = coalesce_gif(gif)
if not frames:
print("No frames found!")
return
# Convert first frame to RGB array, initialize sum array
w, h = frames[0].size
sum_array = np.zeros((h, w, 3), dtype=np.uint16) # 16-bit to avoid overflow
# For each frame:
for f in frames:
# Convert to RGB
rgb = f.convert("RGB")
arr = np.array(rgb, dtype=np.uint16) # shape (H, W, 3)
# Remove near-white by setting it to black
remove_white_make_black(arr, threshold=250)
# Add to sum_array, then clamp to 255
sum_array += arr
np.clip(sum_array, 0, 255, out=sum_array)
# Convert sum_array back to 8-bit
sum_array = sum_array.astype(np.uint8)
# Finally, any pixel that stayed black is presumably "empty," so we set it to white
black_mask = (sum_array[..., 0] == 0) & \
(sum_array[..., 1] == 0) & \
(sum_array[..., 2] == 0)
sum_array[black_mask] = [255, 255, 255]
# Create final Pillow image
final_img = Image.fromarray(sum_array, mode="RGB")
final_img.save("result.png")
print("Done! Wrote result.png.")
if __name__ == "__main__":
main()