Views
No views yet

![]() | ![]() |
![]() | ![]() |
1
2from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
3import torch
4from PIL import Image
5from typing import List, Dict
6import os
7
8def process_vision_info(messages: List[dict]):
9 image_inputs = []
10 video_inputs = []
11 for message in messages:
12 if isinstance(message["content"], list):
13 for item in message["content"]:
14 if item["type"] == "image":
15 image = item["image"]
16 if isinstance(image, str):
17 # Open image with quality improvement
18 image = Image.open(image).convert("RGB")
19 elif isinstance(image, Image.Image):
20 pass
21 else:
22 raise ValueError(f"Unsupported image type: {type(image)}")
23 image_inputs.append(image)
24 elif item["type"] == "video":
25 video_inputs.append(item["video"])
26 return image_inputs if image_inputs else None, video_inputs if video_inputs else None
27
28model_name = "sherif1313/Arabic-handwritten-OCR-4bit-Qwen2.5-VL-3B-v2"
29
30model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
31 model_name,
32 dtype=torch.bfloat16,
33 device_map="auto",
34 trust_remote_code=True
35)
36
37processor = AutoProcessor.from_pretrained(
38 model_name,
39 trust_remote_code=True
40)
41
42def extract_text_from_image(image_path):
43 try:
44 # ✅ Use clearer prompt that requests the complete text
45 messages = [
46 {
47 "role": "user",
48 "content": [
49 {"type": "image", "image": image_path},
50 {"type": "text", "text": "ارجو استخراج النص العربي كاملاً من هذه الصورة من البداية الى النهاية بدون اي اختصار ودون ذيادة او حذف. اقرأ كل المحتوى النصي الموجود في الصورة:"},
51 ],
52 }
53 ]
54
55 # Prepare text and images
56 text = processor.apply_chat_template(
57 messages, tokenize=False, add_generation_prompt=True
58 )
59 image_inputs, video_inputs = process_vision_info(messages)
60
61 # Process inputs with improved settings
62 inputs = processor(
63 text=[text],
64 images=image_inputs,
65 padding=True,
66 return_tensors="pt",
67 ).to(model.device)
68
69 # ✅ Improved generation settings for long texts
70 generated_ids = model.generate(
71 **inputs,
72 max_new_tokens=512, # Significant increase to accommodate long texts
73 min_new_tokens=50, # Minimum to ensure no premature truncation
74 do_sample=False, # For consistent results
75 temperature=0.3, # Balance between creativity and stability
76 top_p=0.9, # For moderate diversity
77 repetition_penalty=1.1, # Prevent repetition
78 pad_token_id=processor.tokenizer.eos_token_id,
79 eos_token_id=processor.tokenizer.eos_token_id,
80 num_return_sequences=1
81 )
82
83 # Extract only the generated text (without user prompt)
84 input_len = inputs.input_ids.shape[1]
85 output_text = processor.batch_decode(
86 generated_ids[:, input_len:],
87 skip_special_tokens=True,
88 clean_up_tokenization_spaces=True # Improve spacing
89 )[0]
90
91 return output_text.strip()
92
93 except Exception as e:
94 return f"Error occurred while processing image: {str(e)}"
95
96def enhance_image_quality(image_path):
97 """Enhance image quality to improve OCR accuracy"""
98 try:
99 img = Image.open(image_path)
100 # Increase resolution if image is small
101 if max(img.size) < 800:
102 new_size = (img.size[0] * 2, img.size[1] * 2)
103 img = img.resize(new_size, Image.Resampling.LANCZOS)
104 return img
105 except:
106 return Image.open(image_path)
107
108if __name__ == "__main__":
109 TEST_IMAGES_DIR = "/media/imges" # Replace with your folder image path
110 IMAGE_EXTENSIONS = ['.png', '.jpg', '.jpeg', '.tif', '.tiff']
111
112 image_files = [
113 os.path.join(TEST_IMAGES_DIR, f)
114 for f in os.listdir(TEST_IMAGES_DIR)
115 if any(f.lower().endswith(ext) for ext in IMAGE_EXTENSIONS)
116 ]
117
118 if not image_files:
119 print("❌ No images found in the folder.")
120 exit()
121
122 print(f"🔍 Found {len(image_files)} images for processing")
123
124 for img_path in sorted(image_files):
125 print(f"\n{'='*50}")
126 print(f"🖼️ Processing: {os.path.basename(img_path)}")
127 print(f"{'='*50}")
128
129 try:
130 # ✅ Use the enhanced function
131 extracted_text = extract_text_from_image(img_path)
132
133 print("📝 Extracted text:")
134 print("-" * 40)
135 print(extracted_text)
136 print("-" * 40)
137
138 # ✅ Calculate text length for comparison
139 text_length = len(extracted_text)
140 print(f"📊 Text length: {text_length} characters")
141
142 except Exception as e:
143 print(f"❌ Error processing {os.path.basename(img_path)}: {e}")