Update app.py

app.py

@@ -65,20 +65,32 @@ def process_zip_files(zip_file_paths):
 def save_captions_to_csv(captions, csv_file_path='image_captions.csv'):
     with open(csv_file_path, mode='w', newline='') as file:
         writer = csv.writer(file)
-        writer.writerow(['Image Name', 'Caption'])
+        writer.writerow(['Image', 'Caption'])
         writer.writerows(captions)
     return csv_file_path
 
-def gr_process(zip_files, image_files):
-    captions = []
-    if zip_files:
-        zip_file_paths = [zip_file.name for zip_file in zip_files]
-        captions.extend(process_zip_files(zip_file_paths))
-    if image_files:
-        image_file_paths = [image_file.name for image_file in image_files]
-        captions.extend(process_images(image_file_paths))
+def gr_process_zip(zip_files):
+    captions = process_zip_files([zip_file.name for zip_file in zip_files])
+    return save_captions_to_csv(captions, 'zip_image_captions.csv')
+
+def gr_process_images(image_files):
+    captions = process_images([image_file.name for image_file in image_files])
+    return save_captions_to_csv(captions, 'image_captions.csv')
+
+def combine_csv_files(file1, file2, output_file='combined_captions.csv'):
+    with open(output_file, mode='w', newline='') as outfile:
+        writer = csv.writer(outfile)
+        writer.writerow(['Image Name', 'Caption'])
+        
+        for file in [file1, file2]:
+            if os.path.exists(file):
+                with open(file, mode='r') as infile:
+                    reader = csv.reader(infile)
+                    next(reader)  # Skip header row
+                    for row in reader:
+                        writer.writerow(row)
     
-    return save_captions_to_csv(captions)
+    return output_file
 
 css = '''
 h1#title {
@@ -113,17 +125,25 @@ with demo:
     with gr.Row():
         with gr.Column(scale=1):
             new_zip_files = gr.File(label="Upload Zip Files", type="filepath", file_count="multiple")
-            zip_captions_btn = gr.Button("Generate Zip Captions")
-            output_zip_file = gr.File(label="Download Caption File")
-        with gr.Column(scale=1):
             new_image_files = gr.File(label="Upload Images", type="filepath", file_count="multiple")
-            image_captions_btn = gr.Button("Generate Image Captions")
-            output_image_file = gr.File(label="Download Caption File")
+        with gr.Column(scale=3):
+            output_zip_file = gr.File(label="Download Zip Captions")
+            output_image_file = gr.File(label="Download Image Captions")
+            combined_file = gr.File(label="Download Combined Captions")
     
-   
+    add_files_btn = gr.Button("Add More Files")
+    generate_zip_captions_btn = gr.Button("Generate Zip Captions")
+    generate_image_captions_btn = gr.Button("Generate Image Captions")
+    combine_files_btn = gr.Button("Combine CSV Files")
 
-    
-    zip_captions_btn.click(fn=gr_process, inputs=[zip_files], outputs=output_zip_file)
-    image_captions_btn.click(fn=gr_process, inputs=[image_files], outputs=output_image_file)
+    # def add_files(zip_files, image_files, new_zip_files, new_image_files):
+    #     zip_files.extend(new_zip_files)
+    #     image_files.extend(new_image_files)
+    #     return zip_files, image_files
+
+    # add_files_btn.click(add_files, inputs=[zip_files, image_files, new_zip_files, new_image_files], outputs=[zip_files, image_files])
+    generate_zip_captions_btn.click(fn=gr_process_zip, inputs=zip_files, outputs=output_zip_file)
+    generate_image_captions_btn.click(fn=gr_process_images, inputs=image_files, outputs=output_image_file)
+    combine_files_btn.click(fn=combine_csv_files, inputs=[output_zip_file, output_image_file], outputs=combined_file)
 
-demo.launch(share=True)
+demo.launch(share=True)