Update app.py

app.py

@@ -1,182 +1,74 @@
 import gradio as gr
-import pandas as pd
 from yt_dlp import YoutubeDL
-from yt_dlp.utils import DownloadError
-from datetime import datetime
-from pathlib import Path
-import plotly.express as px
-import logging
-
-logging.basicConfig(level=logging.INFO)
-
-def _extract_thumbnail_info(video, url, collection_date):
-    """
-    Extracts relevant thumbnail metadata from a yt-dlp entry.
-    """
-    upload_date_raw = video.get('upload_date') or video.get('release_date') or video.get('timestamp')
-    upload_date = pd.to_datetime(upload_date_raw, unit='s', errors='coerce')
-    if isinstance(upload_date_raw, str) and len(upload_date_raw) == 8 and upload_date_raw.isdigit():
-        upload_date = pd.to_datetime(upload_date_raw, format='%Y%m%d', errors='coerce')
-
-    return {
-        'collection_date': collection_date,
-        'profile_url': url,
-        'id': video.get('id'),
-        'title': video.get('title'),
-        'upload_date': upload_date,
-        'uploader': video.get('uploader'),
-        'thumbnail': video.get('thumbnail'),
-        'video_url': video.get('webpage_url') or video.get('url'),
-    }
-
-def collect_thumbnails(profiles_text, profiles_file, cookies_file, max_videos=20, progress=gr.Progress()):
-    """
-    Collects thumbnail URLs and metadata from TikTok profiles using yt-dlp.
-    """
-    urls = []
-    if profiles_text:
-        urls.extend([line.strip() for line in profiles_text.splitlines() if line.strip()])
-    if profiles_file is not None:
-        try:
-            file_content = Path(profiles_file.name).read_text(encoding='utf-8')
-            urls.extend([line.strip() for line in file_content.splitlines() if line.strip()])
-        except Exception as e:
-            return None, None, None, f"❌ Failed to read file: {e}"
-
-    if not urls:
-        return None, None, None, "❌ Please provide at least one TikTok profile URL."
-
-    collection_date = datetime.today().date()
-    all_entries = []
-
+import os
+import tempfile
+import shutil
+
+# -----------------------
+# Helper Functions
+# -----------------------
+def download_tiktok_thumbnails(username_url, cookies_file=None, max_videos=20):
+    tmp_dir = tempfile.mkdtemp()
     ydl_opts = {
-        'quiet': True,
-        'skip_download': True,
-        'extract_flat': False,
-        'http_headers': {
-            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
-                          'AppleWebKit/537.36 (KHTML, like Gecko) '
-                          'Chrome/114.0.0.0 Safari/537.36',
-        },
+        "skip_download": True,
+        "write_thumbnail": True,
+        "outtmpl": os.path.join(tmp_dir, "%(id)s.%(ext)s"),
+        "quiet": True,
     }
-    if cookies_file is not None:
-        ydl_opts['cookiefile'] = cookies_file.name
-
-    with YoutubeDL(ydl_opts) as ydl:
-        for i, url in enumerate(progress.tqdm(urls, desc="Processing Profiles")):
-            try:
-                info = ydl.extract_info(url, download=False)
-                entries = info.get('entries', [])
+    if cookies_file:
+        ydl_opts["cookiefile"] = cookies_file.name
 
-                if entries:
-                    for video in entries[:max_videos]:
-                        if video:
-                            all_entries.append(_extract_thumbnail_info(video, url, collection_date))
-                elif info.get('_type') != 'playlist':
-                    all_entries.append(_extract_thumbnail_info(info, url, collection_date))
-                else:
-                    return None, None, None, f"⚠️ No videos found or profile is private for URL: {url}"
+    try:
+        with YoutubeDL(ydl_opts) as ydl:
+            ydl.download([username_url])
+    except Exception as e:
+        return f"❌ Error: {e}", None, None
 
-            except DownloadError as e:
-                return None, None, None, f"❌ A network error occurred while processing {url}: {e}"
-            except Exception as e:
-                logging.error(f"Error processing {url}: {e}")
-                return None, None, None, f"❌ An unexpected error occurred with {url}: {e}"
+    # Collect all thumbnails
+    thumbs_all = [os.path.join(tmp_dir, f) for f in os.listdir(tmp_dir)
+                  if f.lower().endswith((".jpg", ".webp", ".png"))]
 
-    if not all_entries:
-        return None, None, None, "❌ No thumbnails found."
+    # Apply limiter
+    thumbs = thumbs_all[:max_videos]
 
-    df_thumbs = pd.DataFrame(all_entries)
-    output_file = "tiktok_thumbnails.xlsx"
-    df_thumbs.to_excel(output_file, index=False)
+    if not thumbs:
+        return "❌ No thumbnails found", None, None
 
-    # Collect thumbnail URLs for preview
-    thumb_urls = df_thumbs['thumbnail'].dropna().tolist()
+    return f"✅ Fetched {len(thumbs)} thumbnails", thumbs, tmp_dir
 
-    return output_file, thumb_urls, len(df_thumbs), f"✅ Collected {len(df_thumbs)} thumbnails."
+def prepare_zip(tmp_dir):
+    zip_path = os.path.join(tempfile.gettempdir(), "tiktok_thumbnails.zip")
+    shutil.make_archive(zip_path.replace(".zip",""), 'zip', tmp_dir)
+    return zip_path
 
-def generate_plot(output_file):
-    """
-    Reads the data from the output file and generates a Plotly graph.
-    """
-    if not output_file:
-        return None, "⚠️ No data file to plot."
-    
-    try:
-        df = pd.read_excel(output_file)
-        df['upload_date'] = pd.to_datetime(df['upload_date'], errors='coerce')
-        df = df.dropna(subset=['upload_date', 'uploader'])
-        df['upload_month'] = df['upload_date'].dt.to_period('M').dt.to_timestamp()
+def fetch_and_zip_tiktok(username_url, cookies_file, max_videos):
+    status, thumbs, tmp_dir = download_tiktok_thumbnails(username_url, cookies_file, max_videos)
+    zip_file = prepare_zip(tmp_dir) if thumbs else None
+    return status, thumbs, zip_file
 
-        grouped = df.groupby(['upload_month', 'uploader']).size().reset_index(name='thumbnail_count')
-
-        fig = px.line(
-            grouped,
-            x="upload_month",
-            y="thumbnail_count",
-            color="uploader",
-            markers=True,
-            title="Number of Thumbnails Over Time by Uploader"
-        )
-        fig.update_layout(
-            xaxis_title="Upload Month",
-            yaxis_title="Number of Thumbnails",
-            legend_title="Uploader",
-            hovermode="x unified"
-        )
-        return fig, "✅ Plot generated successfully!"
-    except Exception as e:
-        return None, f"⚠️ Failed to generate plot. Ensure the data file is valid: {e}"
-
-def main_handler(profiles_text, profiles_file, cookies_file, max_videos):
-    """
-    Main function to handle the entire workflow and return Gradio outputs.
-    """
-    # Step 1: Collect thumbnails
-    output_file, thumb_urls, count, collect_msg = collect_thumbnails(
-        profiles_text, profiles_file, cookies_file, max_videos=max_videos
+# -----------------------
+# Gradio Interface
+# -----------------------
+with gr.Blocks() as demo:
+    gr.Markdown("## 🎵 TikTok Account Thumbnails Downloader with Cookies")
+    gr.Markdown(
+        "Upload your `cookies.txt` (exported from your browser) first, then enter a TikTok account URL.\n"
+        "Thumbnails from private or age-restricted videos may require cookies."
     )
-    yield output_file, None, thumb_urls, collect_msg
 
-    # Step 2: Generate plot if collection was successful
-    if output_file:
-        plot, plot_msg = generate_plot(output_file)
-        yield output_file, plot, thumb_urls, plot_msg
-    else:
-        yield None, None, None, collect_msg
+    cookies_upload = gr.File(label="Upload cookies.txt (required for private content)", file_types=[".txt"])
+    url_input = gr.Textbox(label="TikTok Account URL", placeholder="https://www.tiktok.com/@username")
+    max_videos_slider = gr.Slider(minimum=1, maximum=50, step=1, value=20, label="Max Videos to Fetch")
+    fetch_btn = gr.Button("📥 Fetch Thumbnails")
 
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("""
-    # 🖼️ TikTok Profile Thumbnail Collector
-    **Instructions:**
-    - Paste TikTok profile URLs or upload a `.txt` file with one URL per line.
-    - You can upload a `cookies.txt` file to access private profiles.
-    - Adjust the slider to limit the number of thumbnails per profile.
-    - After collection, download the Excel file, preview thumbnails, and view activity over time.
-    ---
-    """)
+    status_output = gr.Textbox(label="Status")
+    thumbs_gallery = gr.Gallery(label="Thumbnails Preview", elem_id="thumbs_gallery", columns=5, height="auto")
+    download_btn = gr.File(label="Download All Thumbnails (ZIP)")
 
-    with gr.Row():
-        profiles_text = gr.Textbox(label="TikTok profile URLs (one per line)", lines=6, placeholder="https://www.tiktok.com/@username")
-        profiles_file = gr.File(label="Or upload a .txt file", file_types=['.txt'])
+    inputs_list = [url_input, cookies_upload, max_videos_slider]
+    outputs_list = [status_output, thumbs_gallery, download_btn]
 
-    cookies_file = gr.File(label="Upload TikTok cookies.txt (optional)", file_types=['.txt'])
-    max_videos_slider = gr.Slider(minimum=1, maximum=50, step=1, value=20, label="Max Videos per Profile")
-    
-    btn = gr.Button("Collect & Analyze Thumbnails")
-    
-    with gr.Row():
-        output_file_display = gr.File(label="Download Excel File")
-        plot_output = gr.Plot(label="Thumbnail Activity Plot")
-    
-    thumbs_gallery = gr.Gallery(label="Thumbnails Preview", columns=5, height="auto")
-    msg = gr.Textbox(label="Status", interactive=False, show_copy_button=True)
+    fetch_btn.click(fetch_and_zip_tiktok, inputs=inputs_list, outputs=outputs_list)
 
-    btn.click(
-        fn=main_handler,
-        inputs=[profiles_text, profiles_file, cookies_file, max_videos_slider],
-        outputs=[output_file_display, plot_output, thumbs_gallery, msg]
-    )
-    
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()