File size: 3,551 Bytes
95ea2c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | import argparse
import json
import os
import cv2
import torch
from PIL import Image
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
from qwen_vl_utils import process_vision_info
MODEL_NAME = "DocShield-7B"
MODEL_PATH = "vankey/DocShield-7B"
IMG_W, IMG_H = 1344, 896
SYSTEM_PROMPT = (
"You are a top-tier image forgery analysis expert, specialized in forensic-level "
"image and text correlation analysis. Based on the input, produce a concise, "
"professional, and accurate forgery analysis report."
)
USER_PROMPT = "Please analyze whether this image is forged and provide an analysis report."
GEN_KWARGS = dict(
do_sample=True,
temperature=1.0,
top_p=1.0,
top_k=0,
repetition_penalty=1.0,
)
def load_model():
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
MODEL_PATH,
torch_dtype=torch.float32,
attn_implementation="eager",
device_map="auto",
)
model.eval()
return model
def load_image(image_path):
image = cv2.imread(image_path)
if image is None:
raise FileNotFoundError(f"Cannot read image: {image_path}")
image = cv2.resize(image, (IMG_W, IMG_H))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
return Image.fromarray(image)
def build_messages(image_path):
img = load_image(image_path)
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{
"role": "user",
"content": [
{"type": "image", "image": img},
{"type": "text", "text": USER_PROMPT},
],
},
]
return messages
def infer(model, processor, image_path, max_new_tokens=8192):
messages = build_messages(image_path)
text = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
text=[text],
images=image_inputs,
videos=video_inputs,
padding=True,
return_tensors="pt",
).to(model.device)
with torch.no_grad():
output_ids = model.generate(**inputs, **GEN_KWARGS, max_new_tokens=max_new_tokens)
generated_ids = output_ids[:, inputs["input_ids"].shape[1]:]
result = processor.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)[0]
return result.strip()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--image", required=True)
parser.add_argument("--output", default=None)
parser.add_argument("--max-new-tokens", type=int, default=8192)
args = parser.parse_args()
if not os.path.exists(args.image):
raise FileNotFoundError(f"Image not found: {args.image}")
print(f"[{MODEL_NAME}] loading model: {MODEL_PATH}")
processor = AutoProcessor.from_pretrained(MODEL_PATH)
model = load_model()
print(f"[{MODEL_NAME}] inferring image: {args.image}")
answer = infer(model, processor, args.image, max_new_tokens=args.max_new_tokens)
output_data = {"model": MODEL_NAME, "image": args.image, "answer": answer}
output_path = args.output or os.path.join(
os.getcwd(), f"{os.path.splitext(os.path.basename(args.image))[0]}.json"
)
with open(output_path, "w", encoding="utf-8") as f:
json.dump(output_data, f, ensure_ascii=False, indent=4)
print("\n--- result ---")
print(answer)
print(f"\n--- saved to: {output_path} ---")
if __name__ == "__main__":
main()
|