Upload 2 files

app.py

@@ -0,0 +1,98 @@
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, ViTImageProcessor, VisionEncoderDecoderModel
+import zipfile
+import os
+import csv
+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")
+
+def predict(image, max_length=64, num_beams=4):
+    # Process the input image
+    image = image.convert('RGB')
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
+    
+    # Generate the caption
+    caption_ids = model.generate(pixel_values, max_length=max_length, num_beams=num_beams)[0]
+    
+    # Decode and clean the generated caption
+    caption = tokenizer.decode(caption_ids, skip_special_tokens=True)
+    return caption
+
+def process_zip_file(zip_file_path):
+    # 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}")
+    
+    # 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 gr_process_zip(zip_file):
+    zip_file_path = zip_file.name
+    return process_zip_file(zip_file_path)
+
+css = '''
+h1#title {
+  text-align: center;
+}
+h3#header {
+  text-align: center;
+}
+img#overview {
+  max-width: 800px;
+  max-height: 600px;
+}
+img#style-image {
+  max-width: 1000px;
+  max-height: 600px;
+}
+.gr-image {
+  max-width: 150px;  /* Set a small box for the image */
+  max-height: 150px;
+}
+'''
+
+demo = gr.Blocks(css=css)
+
+with demo:
+    gr.Markdown('''<h1 id="title">Image Caption 🖼️</h1>''')
+    gr.Markdown('''Made by : No. Fa.''')
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_zip = gr.File(label="Upload your Zip File", type="file")
+        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)
+
+demo.launch()

requirements.txt

@@ -0,0 +1,4 @@
+torch
+transformers 
+gradio 
+pillow