import gradio as gr
from huggingface_hub import hf_hub_download, HfFolder
from PIL import Image
import requests, torch
import numpy as np
from io import BytesIO
import plotly.graph_objects as go
import os

def load_ScanNet_sample(data_path):
    
    all_data = torch.load(data_path)
    
    point = np.array(all_data['coord'])
    color = np.array(all_data['color'])
 
    point = point - point.min(axis=0)
    point = point / point.max(axis=0)
    color = color / 255.
    return point, color

def show_logo():
    repo_id = "ZiyuG/Cache"
    filename = "scene0000_00.pth"
    token = os.getenv('HF_TOKEN')
    print("token:", token)

    try:
        file_path = hf_hub_download(repo_id=repo_id, filename=filename, use_auth_token=token, repo_type='dataset')
        point, color = load_ScanNet_sample(file_path)
        if point.shape[0] > 100000:
            indices = np.random.choice(point.shape[0], 100000, replace=False)
            point = point[indices]
            color = color[indices]
    except Exception as e:
        print(e)
        point = np.random.rand(8000, 3)
        color = np.random.rand(8000, 3)
    
    fig = go.Figure(
        data=[
            go.Scatter3d(
                x=point[:,0], y=point[:,1], z=point[:,2],
                mode='markers',
                marker=dict(size=1, color=color, opacity=0.5),
            )
        ],
        layout=dict(
            scene=dict(
                xaxis=dict(visible=False),
                yaxis=dict(visible=False),
                zaxis=dict(visible=False),
                aspectratio=dict(x=1, y=1, z=1), 
                camera=dict(eye=dict(x=1.5, y=1.5, z=1.5))
            )
        )
    )
    return fig

iface = gr.Interface(fn=show_logo, inputs=[], outputs=gr.Plot(), title="Display Logo")
iface.launch(share=True)