Update run.py

run.py

@@ -9,7 +9,7 @@ from huggingface_hub import hf_hub_download, snapshot_download
 import pyarrow.parquet as pq
 import hashlib
 import os
-
+import csv
 
 def load_hf_dataset(dataset_path, auth_token):
     dataset = load_dataset(dataset_path, token=auth_token)
@@ -18,47 +18,53 @@ def load_hf_dataset(dataset_path, auth_token):
 
 def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, endpoint, video_src, video_hf, video_hf_auth, parquet_index, video_od, video_od_auth, video_gd, video_gd_auth, frame_format, frame_limit):
     progress_info = []
-    if video_src:
-        video = video_src
-        processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
-        frames = processor._decode(video)
-        base64_list = processor.to_base64_list(frames)
-        debug_image = processor.concatenate(frames)
-        if not key or not endpoint:
-            return "", f"API key or endpoint is missing. Processed {len(frames)} frames.", debug_image
-        api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
-        caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
-        progress_info.append(f"Using model '{model}' with {len(frames)} frames extracted.")
-        return f"{caption}", "\n".join(progress_info), debug_image
-    elif video_hf and video_hf_auth:
-        current_file_path = os.path.abspath(__file__)
-        current_directory = os.path.dirname(current_file_path)
-        progress_info.append('Begin processing Hugging Face dataset.')
-        all_captions = []
-        temp_parquet_file = hf_hub_download(
-            repo_id=video_hf,
-            filename='data/' + str(parquet_index).zfill(6) + '.parquet',
-            repo_type="dataset",
-            token=video_hf_auth,
-        )
-        parquet_file = pq.ParquetFile(temp_parquet_file)
-        for batch in parquet_file.iter_batches(batch_size=1):
-            df = batch.to_pandas()
-            video = df['video'][0]
-            md5 = hashlib.md5(video).hexdigest()
-            with tempfile.NamedTemporaryFile(dir=current_directory) as temp_file:
-                temp_file.write(video)
-                video_path = temp_file.name
-                processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
-                frames = processor._decode(video_path)
-                base64_list = processor.to_base64_list(frames)
-                api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
-                caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
-                all_captions.append(caption)
-                progress_info.append(f"Processed video with MD5: {md5}")
-        return "\n\n\n".join(all_captions), "\n".join(progress_info), None
-    else:
-        return "", "No video source selected.", None
+    csv_filename = "captions.csv"
+    with open(csv_filename, mode='w', newline='') as csv_file:
+        fieldnames = ['md5', 'caption']
+        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
+        writer.writeheader()
+
+        if video_src:
+            video = video_src
+            processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
+            frames = processor._decode(video)
+            base64_list = processor.to_base64_list(frames)
+            debug_image = processor.concatenate(frames)
+            if not key or not endpoint:
+                return "", f"API key or endpoint is missing. Processed {len(frames)} frames.", debug_image
+            api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
+            caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
+            progress_info.append(f"Using model '{model}' with {len(frames)} frames extracted.")
+            writer.writerow({'md5': 'single_video', 'caption': caption})
+            return f"{caption}", "\n".join(progress_info), debug_image
+        elif video_hf and video_hf_auth:
+            current_file_path = os.path.abspath(__file__)
+            current_directory = os.path.dirname(current_file_path)
+            progress_info.append('Begin processing Hugging Face dataset.')
+            temp_parquet_file = hf_hub_download(
+                repo_id=video_hf,
+                filename='data/' + str(parquet_index).zfill(6) + '.parquet',
+                repo_type="dataset",
+                token=video_hf_auth,
+            )
+            parquet_file = pq.ParquetFile(temp_parquet_file)
+            for batch in parquet_file.iter_batches(batch_size=1):
+                df = batch.to_pandas()
+                video = df['video'][0]
+                md5 = hashlib.md5(video).hexdigest()
+                with tempfile.NamedTemporaryFile(dir=current_directory) as temp_file:
+                    temp_file.write(video)
+                    video_path = temp_file.name
+                    processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
+                    frames = processor._decode(video_path)
+                    base64_list = processor.to_base64_list(frames)
+                    api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
+                    caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
+                    writer.writerow({'md5': md5, 'caption': caption})
+                    progress_info.append(f"Processed video with MD5: {md5}")
+            return csv_filename, "\n".join(progress_info), None
+        else:
+            return "", "No video source selected.", None
 
 with gr.Blocks() as Core:
     with gr.Row(variant="panel"):
@@ -124,12 +130,12 @@ with gr.Blocks() as Core:
                             video_gd = gr.Text()
                             video_gd_auth = gr.Text(label="Google Drive Access Token")
                 caption_button = gr.Button("Caption", variant="primary", size="lg")
+                csv_link = gr.File(label="Download CSV", interactive=False)
         caption_button.click(
             fast_caption, 
             inputs=[sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, endpoint, video_src, video_hf, video_hf_auth, parquet_index, video_od, video_od_auth, video_gd, video_gd_auth, frame_format, frame_limit], 
-            outputs=[result, info, frame]
+            outputs=[csv_link, info, frame]
         )
 
 if __name__ == "__main__":
-    Core.launch()
-
+    Core.launch()