Upload layer_diff_dataset/tagger_jsh2.py with huggingface_hub

layer_diff_dataset/tagger_jsh2.py

@@ -0,0 +1,79 @@
+from modelscope import (
+    snapshot_download, AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+)
+import torch
+import os
+import json
+# model_id = 'qwen/Qwen-VL-Chat'
+# revision = 'v1.1.0'
+
+# model_dir = snapshot_download(model_id, revision=revision)
+seed = 1234
+torch.manual_seed(seed)
+model_dir = "/mnt/workspace/workgroup/sihui.jsh/VITON-HD/diffusers/qwen/Qwen-VL-Chat"
+
+# 请注意：分词器默认行为已更改为默认关闭特殊token攻击防护。
+tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+# 打开bf16精度，A100、H100、RTX3060、RTX3070等显卡建议启用以节省显存 balanced
+model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="sequential", trust_remote_code=True, fp16=True).eval()
+# 打开fp16精度，V100、P100、T4等显卡建议启用以节省显存
+# model = AutoModelForCausalLM.from_pretrained(model_dir, trust_remote_code=True, fp16=True).eval()
+# 使用CPU进行推理，需要约32GB内存
+# model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="cpu", trust_remote_code=True).eval()
+# 默认使用自动模式，根据设备自动选择精度
+# model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True).eval()
+
+# 可指定不同的生成长度、top_p等相关超参
+model.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)
+
+# 第一轮对话 1st dialogue turn
+
+# image_dir_path = '/home/nfs/wyy/data/deepfashion/Data/test_lst_512_png'
+# image_dir_path = '/home/nfs/wyy/data/VTO/VITON-HD-512/train/image'
+# image_dir_path = '/mnt/workspace/workgroup/sihui.jsh/VITON-HD/try/image'
+root_dir = '/mnt/workspace/workgroup/sihui.jsh/layer_diff_dataset/train'
+image_dir_path = os.path.join(root_dir,'im')
+
+from tqdm import tqdm
+# os.makedirs('./test_cap', exist_ok=True)
+image_dir = os.listdir(image_dir_path)
+image_dir.sort()
+# sort
+# image_dir = sorted(image_dir, key=lambda x: x.split('.')[0])
+import numpy as np
+
+# test_dir = image_dir_path.replace('cloth_resize_combined1','caption_new1')
+# os.makedirs(test_dir, exist_ok=True)
+
+# prompt = 'Describe the image'
+prompt = 'Only describe the background of the image.'
+with open(os.path.join(root_dir,'im_rgba.json'), 'r') as file:
+    data = json.load(file)
+# part_size = len(image_dir) // 2
+pbar = tqdm(enumerate(image_dir),total=len(image_dir))
+for idx, image in pbar:
+    # if cnt<1000:
+    #     continue
+    # if cnt==2500:
+    #     break
+    image_path = os.path.join(image_dir_path, image)
+    
+    query = tokenizer.from_list_format([
+        {'image': image_path},
+        {'text': prompt},
+    ])
+    response, _ = model.chat(tokenizer, query=query, history=None)
+    if data[idx]["images"]==image_path:
+        data[idx]["bg_prompt"] = response
+    else:
+        print('error!')
+        print(image_path)
+        print(data[idx]["images"])
+        break
+    # print(response)
+    # shutil.copy(image_path, os.path.join(test_dir, image))   
+    # print(os.path.join(test_dir,'image2text.txt')) 
+# with open(os.path.join(root_dir,'.txt'), 'a') as f:
+#     f.write(image + '\t' + response + '\n')
+with open(os.path.join(root_dir,'im_rgba_new.json'), 'w') as file:
+    json.dump(data, file, indent=4)