Upload new_captioner.py

new_captioner.py

@@ -0,0 +1,109 @@
+import base64
+import requests
+import os
+from openai import OpenAI
+from tqdm import tqdm
+import time
+import sys
+
+# Проверка наличия аргумента командной строки
+if len(sys.argv) < 2:
+    print("Please, provide the path to image folder.")
+    sys.exit(1)
+
+# Get the path to image dir from command line.
+image_dir = sys.argv[1]
+
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:8000/v1"
+client = OpenAI(
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+
+model_type = client.models.list().data[0].id
+print(f'model_type: {model_type}')
+
+# Function to encode the image
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')
+
+# Directories
+#dir with tags captions from wd tagger
+txt_dir = './txt/'
+#dir with result captions
+maintxt_dir = './maintxt/'
+image_path =''
+
+# Ensure the output directory exists
+os.makedirs(maintxt_dir, exist_ok=True)
+
+# Get list of all JPEG images in the directory
+image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.jpeg'))]
+
+total_files = len(image_files)
+start_time = time.time()
+
+progress_bar = tqdm(total=total_files, unit='file', bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]')
+total_elapsed_time = 0
+processed_files = 0
+
+# Process all images in the image directory
+for image_file in image_files:
+    image_path = os.path.join(image_dir, image_file)
+    txt_file = os.path.join(txt_dir, os.path.splitext(image_file)[0] + '.txt')
+    output_file = os.path.join(maintxt_dir, os.path.splitext(image_file)[0] + '.txt')
+
+    # Read tags from the corresponding txt file
+    with open(txt_file, 'r') as f:
+        tags = f.read().strip()
+
+    base64_image = encode_image(image_path)
+
+    step_start_time = time.time()
+
+    chat_response = client.chat.completions.create(
+        model="./phi3_v14_800-merged",
+        messages=[{
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"Make a caption that describe this image. Here is the tags for this image: {tags}"},
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/jpeg;base64,{base64_image}"
+                    },
+                },
+            ],
+        }],
+        extra_body={'repetition_penalty': 1.05, 'top_k': -1,'top_p': 1,'temperature': 0, 'use_beam_search': True, 'best_of':5},
+    )
+
+    step_end_time = time.time()
+    step_time = step_end_time - step_start_time
+    total_elapsed_time += step_time
+    remaining_time = (total_elapsed_time / (processed_files + 1)) * (total_files - processed_files - 1)
+
+    # Convert remaining time to hours, minutes and seconds
+    remaining_hours = int(remaining_time // 3600)
+    remaining_minutes = int((remaining_time % 3600) // 60)
+    remaining_seconds = int(remaining_time % 60)
+
+    # Extract the content from the response
+    content = chat_response.choices[0].message.content
+    content = content.lstrip()
+    # Write the content to the output file
+    with open(output_file, 'w', encoding='utf-8') as f:
+        f.write(content)
+
+    print(f"\n\nFile {image_file}\nProcessing time: {step_time:.2f} seconds\n{content}")
+    print(f"Response saved to file: {output_file}")
+    
+    processed_files += 1
+    progress_bar.update(1)
+    progress_bar.set_postfix(remaining=f'{remaining_hours:02d}:{remaining_minutes:02d}:{remaining_seconds:02d}', refresh=True)
+
+progress_bar.close()
+print("All images processed.")
+print(f"Total time: {time.time() - start_time:.2f} seconds")