Upload 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
@@ -25,29 +25,50 @@ def predict(image, max_length=64, num_beams=4):
     caption = tokenizer.decode(caption_ids, skip_special_tokens=True)
     return caption
 
-def process_zip_file(zip_file_path):
+def process_images(image_files):
+    captions = []
+    for image_file in image_files:
+        try:
+            # Open and verify the image
+            with Image.open(image_file) as img:
+                caption = predict(img)
+                captions.append((os.path.basename(image_file), caption))
+        except Exception as e:
+            print(f"Skipping file {image_file}: {e}")
+    
+    # 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 Name', 'Caption'])
+        writer.writerows(captions)
+    
+    return csv_file_path
+
+def process_zip_files(zip_file_paths):
     # Create a directory to extract images
     extract_dir = 'extracted_images'
     os.makedirs(extract_dir, exist_ok=True)
-
-    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
-        zip_ref.extractall(extract_dir)
     
-    # Verify extracted files and process images
     captions = []
-    for root, dirs, files in os.walk(extract_dir):
-        for file in files:
-            file_path = os.path.join(root, file)
-            try:
-                # Open and verify the image
-                with Image.open(file_path) as img:
-                    caption = predict(img)
-                    captions.append((file, caption))
-            except Exception as e:
-                print(f"Skipping file {file}: {e}")
+    for zip_file_path in zip_file_paths:
+        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+            zip_ref.extractall(extract_dir)
+        
+        # Verify extracted files and process images
+        for root, dirs, files in os.walk(extract_dir):
+            for file in files:
+                file_path = os.path.join(root, file)
+                try:
+                    # Open and verify the image
+                    with Image.open(file_path) as img:
+                        caption = predict(img)
+                        captions.append((file, caption))
+                except Exception as e:
+                    print(f"Skipping file {file}: {e}")
     
     # Save the results to a CSV file
-    csv_file_path = 'image_captions.csv'
+    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'])
@@ -55,9 +76,15 @@ def process_zip_file(zip_file_path):
     
     return csv_file_path
 
-def gr_process_zip(zip_file):
-    zip_file_path = zip_file.name
-    return process_zip_file(zip_file_path)
+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
 
 css = '''
 h1#title {
@@ -88,11 +115,12 @@ with demo:
     
     with gr.Row():
         with gr.Column(scale=1):
-            input_zip = gr.File(label="Upload your Zip File", type="file")
+            input_zip = gr.File(label="Upload Zip Files", type="filepath", file_count="multiple")
+            input_images = gr.File(label="Upload Images", type="filepath", file_count="multiple")
         with gr.Column(scale=3):
             output_file = gr.File(label="Download Caption File")
         
     btn = gr.Button("Generate Captions")
-    btn.click(fn=gr_process_zip, inputs=input_zip, outputs=output_file)
+    btn.click(fn=gr_process, inputs=[input_zip, input_images], outputs=output_file)
 
 demo.launch()