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-01").to(device)
-feature_extractor = ViTImageProcessor.from_pretrained("NourFakih/Vit-GPT2-COCO2017Flickr-01")
-tokenizer = AutoTokenizer.from_pretrained("NourFakih/Vit-GPT2-COCO2017Flickr-01")
+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")
 
 def predict(image, max_length=64, num_beams=4):
     # Process the input image
@@ -36,7 +36,14 @@ def process_images(image_files):
         except Exception as e:
             print(f"Skipping file {image_file}: {e}")
     
-    return captions
+    # Save the results to a CSV file
+    csv_file_path = 'image_captions.csv'
+    with open(csv_file_path, mode='w', newline='') as file:
+        writer = csv.writer(file)
+        writer.writerow(['Image', 'Caption'])
+        writer.writerows(captions)
+    
+    return csv_file_path
 
 def process_zip_files(zip_file_paths):
     # Create a directory to extract images
@@ -60,23 +67,24 @@ def process_zip_files(zip_file_paths):
                 except Exception as e:
                     print(f"Skipping file {file}: {e}")
     
-    return captions
-
-def save_captions_to_csv(captions, csv_file_path='image_captions.csv'):
+    # Save the results to a CSV file
+    csv_file_path = 'zip_image_captions.csv'
     with open(csv_file_path, mode='w', newline='') as file:
         writer = csv.writer(file)
         writer.writerow(['Image Name', 'Caption'])
         writer.writerows(captions)
+    
     return csv_file_path
 
-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 gr_process(zip_files, image_files):
+    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:
+        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)
@@ -135,8 +143,8 @@ with demo:
     combine_files_btn = gr.Button("Combine CSV 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)
+    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)
 
 demo.launch()