Update app.py

app.py

@@ -9,9 +9,9 @@ from PIL import Image
 device = 'cpu'
 
 # Load the pretrained model, feature extractor, and tokenizer
-model = VisionEncoderDecoderModel.from_pretrained("NourFakih/Vit-GPT2-COCO2017Flickr-02").to(device)
-feature_extractor = ViTImageProcessor.from_pretrained("NourFakih/Vit-GPT2-COCO2017Flickr-02")
-tokenizer = AutoTokenizer.from_pretrained("NourFakih/Vit-GPT2-COCO2017Flickr-02")
+model = VisionEncoderDecoderModel.from_pretrained("NourFakih/Vit-GPT2-COCO2017Flickr-01").to(device)
+feature_extractor = ViTImageProcessor.from_pretrained("NourFakih/Vit-GPT2-COCO2017Flickr-01")
+tokenizer = AutoTokenizer.from_pretrained("NourFakih/Vit-GPT2-COCO2017Flickr-01")
 
 def predict(image, max_length=64, num_beams=4):
     # Process the input image
@@ -35,10 +35,9 @@ def process_images(image_files):
                 captions.append((os.path.basename(image_file), caption))
         except Exception as e:
             print(f"Skipping file {image_file}: {e}")
-    csv_file_path = save_captions_to_csv(captions, csv_file_path='zip_captions.csv')
-    return csv_file_path
+    
+    return captions
 
-   
 def process_zip_files(zip_file_paths):
     # Create a directory to extract images
     extract_dir = 'extracted_images'
@@ -60,8 +59,8 @@ def process_zip_files(zip_file_paths):
                         captions.append((file, caption))
                 except Exception as e:
                     print(f"Skipping file {file}: {e}")
-        csv_file_path = save_captions_to_csv(captions, csv_file_path='zip_captions.csv')
-    return csv_file_path
+    
+    return captions
 
 def save_captions_to_csv(captions, csv_file_path='image_captions.csv'):
     with open(csv_file_path, mode='w', newline='') as file:
@@ -71,34 +70,15 @@ def save_captions_to_csv(captions, csv_file_path='image_captions.csv'):
     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]
-        return process_zip_files(zip_file_paths)
-    elif image_files:
+        captions.extend(process_zip_files(zip_file_paths))
+    if image_files:
         image_file_paths = [image_file.name for image_file in image_files]
-        return process_images(image_file_paths)
-    else:
-        return None
-
-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)
+        captions.extend(process_images(image_file_paths))
     
-    return output_file
-
-def add_files(zip_files, image_files, new_zip_files, new_image_files):
-    zip_files += new_zip_files
-    image_files += new_image_files
-    return zip_files, image_files
+    return save_captions_to_csv(captions)
 
 css = '''
 h1#title {
@@ -125,7 +105,7 @@ demo = gr.Blocks(css=css)
 
 with demo:
     gr.Markdown('''<h1 id="title">Image Caption 🖼️</h1>''')
-    gr.Markdown('''Made by : N. F.''')
+    gr.Markdown('''Made by : No. Fa.''')
     
     zip_files = gr.State([])
     image_files = gr.State([])
@@ -133,21 +113,19 @@ with demo:
     with gr.Row():
         with gr.Column(scale=1):
             new_zip_files = gr.File(label="Upload Zip Files", type="filepath", file_count="multiple")
-            generate_zip_captions_btn = gr.Button("Generate Zip Captions")
-            output_zip_file = gr.File(label="Download Zip Captions", interactive=False)
-        with gr.Column(scale=1):
             new_image_files = gr.File(label="Upload Images", type="filepath", file_count="multiple")
-            generate_image_captions_btn = gr.Button("Generate Image Captions")
-            output_image_file = gr.File(label="Download Image Captions", interactive=False)
+        with gr.Column(scale=3):
+            output_file = gr.File(label="Download Caption File")
+    
+    add_files_btn = gr.Button("Add More Files")
+    generate_captions_btn = gr.Button("Generate Captions")
+
+    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
 
-    with gr.Row():
-        add_files_btn = gr.Button("Add More Files")
-        combine_files_btn = gr.Button("Combine CSV Files")
-        combined_file = gr.File(label="Download Combined Captions", interactive=False)
-      
     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, inputs=zip_files, outputs=output_zip_file)
-    generate_image_captions_btn.click(fn=gr_process, 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)
+    generate_captions_btn.click(fn=gr_process, inputs=[zip_files, image_files], outputs=output_file)
 
 demo.launch(share=True)