Skip to content

Instantly share code, notes, and snippets.

@eebmagic
Created February 1, 2024 03:41
Show Gist options
  • Save eebmagic/95f6def6e02d46fae79fe543f4213086 to your computer and use it in GitHub Desktop.
Save eebmagic/95f6def6e02d46fae79fe543f4213086 to your computer and use it in GitHub Desktop.
Chroma to Plotly 3D scatter template

Reduce and Plot

Reduce data from a collection to 3D datapoints with PCA reduction.

Then use plotly to plot to a 3d scatter plot.

import plotly.express as px
from sklearn.decomposition import PCA

pca = PCA(n_components=3)
pcaResult = pca.fit_transform(mat)
print(pcaResult.shape)

df = pd.DataFrame(pcaResult)
df = df.assign(label=[str(l) for l in labels], ids=ids)
df['labelInt'] = pd.to_numeric(df['label'])
df['date'] = [m['datestr'] for m in metas]
df['filename'] = [m['filename'] for m in metas]
df['project'] = [m['project'] for m in metas]

subset = df.query('labelInt >= 0')
# subset = df.query('labelInt >= -1')
subset = subset.sort_values(by="labelInt")

fig = px.scatter_3d(
    subset,
    x=0, y=1, z=2,
    color='label',
    hover_data=['project', 'filename', 'date'],
    title=f"{collection.name} ({dims} dims)"
)
fig.update_scenes(
    xaxis_range=[subset[0].min(), subset[0].max()],
    yaxis_range=[subset[1].min(), subset[1].max()],
    zaxis_range=[subset[2].min(), subset[2].max()],
    aspectmode='cube'
)

# print(subset)
print(subset['labelInt'].value_counts().sort_index())
fig.show(renderer="browser") 

Make animation

Run this command to convert all the frame images to a video:

ffmpeg -framerate 30 -i frame_%d.png -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p output.mp4
import plotly.graph_objects as go


# Your existing plot creation code
fig = px.scatter_3d(
    subset,
    x=0, y=1, z=2,
    color='label',
    hover_data=['project', 'filename', 'date'],
    title=f"{collection.name} ({dims} dims)"
)

# Increasing the number of frames and decreasing the frame duration
n_frames = 300  # Increase the number of frames for smoother animation
# n_frames = 50  # Increase the number of frames for smoother animation
# frame_duration = 20  # Decrease duration (in milliseconds) for faster animation
frame_duration = 80  # Decrease duration (in milliseconds) for faster animation
width = 1920
height = 1080

frames = []
for i, t in enumerate(np.linspace(0, 2 * np.pi, n_frames)):  # Complete rotation in n_frames steps
    frame = go.Frame(
        layout=dict(scene_camera=dict(eye=dict(x=2 * np.sin(t), y=2 * np.cos(t), z=0.1)))
    )
    frames.append(frame)
    if (i // n_frames) % 10 == 0:
        print(f"Writing frame: {t / (2 * np.pi)}")
    fig.update_layout(
        scene_camera=dict(eye=dict(x=2 * np.sin(t), y=2 * np.cos(t), z=0.1))
    )
    fig.write_image(
        f"frames/frame_{i}.png",
        width=width,
        height=height,
    )
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment