SSD-DIS: A Semi-Synthetic Shadow Dataset for Document Images.

1. import os
2. import random
3. from PIL import Image
4. import torchvision.transforms as transforms
5. from torchvision.io import read_image
7. # ============================================================================
8. # CONFIGURATION PARAMETERS - USERS CAN MODIFY THESE PATHS AND PARAMETERS
9. # ============================================================================
11. # Shadow mask library folder (contains all available shadow mask images)
12. SHADOW_MASK_LIBRARY_DIR = 'D:\\shadow_document_datasets\\semi-synthetic dataset\\shadows'
14. # Shadow-free document images folder (users need to provide their own shadow-free document images)
15. # Note: Choose the appropriate folder based on your needs (train, val, or test)
16. SHADOW_FREE_IMAGES_DIR = 'D:\\shadow_document_datasets\\semi-synthetic dataset\\shadow_free_images\\test' # Change to 'train' or 'val' as needed
18. # Output folder settings (it's recommended to create these directories before running)
19. OUTPUT_SHADOW_MASKS_DIR = 'D:\\shadow_document_datasets\\semi-synthetic dataset\\vers_shadow_masks\\val'
20. OUTPUT_SHADOW_IMAGES_DIR = 'D:\\shadow_document_datasets\\semi-synthetic dataset\\vers_shadow_images\\val'
21. OUTPUT_ADJUSTED_MASKS_DIR = 'D:\\shadow_document_datasets\\semi-synthetic dataset\\vers_shadow_masks_intensity_adjusted\\val'
23. # Shadow intensity range (calibrated based on human perception, recommended: 0.15-0.8)
24. # The intensity range used in the paper is (0.15, 0.8)
25. SHADOW_INTENSITY_RANGE = (0.15, 0.8) # Can be adjusted to (0.8, 1.0) for stronger shadows
27. # ============================================================================
28. # MAIN PROGRAM STARTS HERE
29. # ============================================================================
31. def synthesize_shadow_dataset():
32. """
33. Main function: Batch synthesize document shadow dataset
34. Functionality: Randomly select masks from the shadow mask library, resize them,
35. and apply them to shadow-free document images to generate
36. document images with synthetic shadows and corresponding shadow masks.
37. """
39. # 1. Check if input directories exist
40. if not os.path.exists(SHADOW_FREE_IMAGES_DIR):
41. print(f"Error: Shadow-free images directory does not exist - {SHADOW_FREE_IMAGES_DIR}")
42. return
44. if not os.path.exists(SHADOW_MASK_LIBRARY_DIR):
45. print(f"Error: Shadow mask library directory does not exist - {SHADOW_MASK_LIBRARY_DIR}")
46. return
48. # 2. Create output directories if they don't exist
49. for output_dir in [OUTPUT_SHADOW_MASKS_DIR, OUTPUT_SHADOW_IMAGES_DIR, OUTPUT_ADJUSTED_MASKS_DIR]:
50. os.makedirs(output_dir, exist_ok=True)
52. # 3. Get all shadow-free document image files
53. shadow_free_files = [f for f in os.listdir(SHADOW_FREE_IMAGES_DIR)
54. if os.path.isfile(os.path.join(SHADOW_FREE_IMAGES_DIR, f))]
56. print(f"Starting to process {len(shadow_free_files)} shadow-free document images...")
57. print(f"Shadow intensity range: {SHADOW_INTENSITY_RANGE[0]:.2f} - {SHADOW_INTENSITY_RANGE[1]:.2f}")
58. print(f"Output directory: {OUTPUT_SHADOW_IMAGES_DIR}")
60. # 4. Get all mask files from the shadow mask library
61. shadow_mask_files = os.listdir(SHADOW_MASK_LIBRARY_DIR)
62. if not shadow_mask_files:
63. print("Error: No image files found in the shadow mask library")
64. return
66. # 5. Process each shadow-free document image
67. for shadow_free_file in shadow_free_files:
68. # 5.1 Randomly select a shadow mask from the library
69. random_mask_file = random.choice(shadow_mask_files)
71. # 5.2 Read the original shadow mask and resize it to match the document image
72. mask_source_path = os.path.join(SHADOW_MASK_LIBRARY_DIR, random_mask_file)
73. mask_target_path = os.path.join(OUTPUT_SHADOW_MASKS_DIR, shadow_free_file)
75. # Open the shadow mask and resize it
76. with Image.open(mask_source_path) as mask_img:
77. # Read the corresponding shadow-free document image to get the target size
78. shadow_free_path = os.path.join(SHADOW_FREE_IMAGES_DIR, shadow_free_file)
79. with Image.open(shadow_free_path) as doc_img:
80. # Resize the shadow mask to the same dimensions as the document image
81. resized_mask = mask_img.resize(doc_img.size, Image.Resampling.LANCZOS)
82. # Save the resized mask
83. resized_mask.save(mask_target_path)
85. print(f"Processing: {shadow_free_file} | Using mask: {random_mask_file} | Saved to: {mask_target_path}")
87. # 5.3 Synthesize the shadowed document image
88. synthesize_shadow_image(shadow_free_file)
90. print("=" * 80)
91. print("Data processing completed!")
92. print(f"Number of generated shadowed images: {len(shadow_free_files)}")
93. print(f"Output file locations:")
94. print(f" - Shadow masks: {OUTPUT_SHADOW_MASKS_DIR}")
95. print(f" - Synthetic shadow images: {OUTPUT_SHADOW_IMAGES_DIR}")
96. print(f" - Adjusted masks: {OUTPUT_ADJUSTED_MASKS_DIR}")
98. def synthesize_shadow_image(shadow_free_filename):
99. """
100. Synthesize a single shadowed document image
102. Parameters:
103. shadow_free_filename: Filename of the shadow-free document image
104. """
105. # 1. Read the shadow-free document image
106. shadow_free_path = os.path.join(SHADOW_FREE_IMAGES_DIR, shadow_free_filename)
107. shadow_free_image = read_image(shadow_free_path)
109. # 2. Read the resized shadow mask
110. shadow_mask_path = os.path.join(OUTPUT_SHADOW_MASKS_DIR, shadow_free_filename)
111. shadow_mask = read_image(shadow_mask_path)
113. # 3. Generate random shadow intensity for the current image
114. # Note: Using the globally defined SHADOW_INTENSITY_RANGE
115. shadow_intensity = random.uniform(SHADOW_INTENSITY_RANGE[0], SHADOW_INTENSITY_RANGE[1])
117. # 4. Convert the shadow mask to a tensor and adjust intensity
118. # Shadow masks are typically binary images (0-255), normalize to 0-1 range
119. shadow_tensor = 1.0 - shadow_mask.float() / 255.0
121. # Apply shadow intensity: shadowed area = shadow_intensity, non-shadowed area = 1.0
122. shadow_tensor = shadow_intensity + (1.0 - shadow_intensity) * shadow_tensor
124. # 5. Normalize the shadow-free document image to [0, 1] range
125. shadow_free_tensor = shadow_free_image.float() / 255.0
127. # 6. Expand the channel dimension of the shadow tensor to match the document image (single-channel -> three-channel)
128. shadow_tensor = shadow_tensor.expand_as(shadow_free_tensor)
130. # 7. Apply shadow: Multiply the document image with the shadow tensor
131. # Shadowed areas become darker, non-shadowed areas remain unchanged
132. shadowed_image = shadow_free_tensor * shadow_tensor
134. # 8. Save the synthesized shadowed document image
135. shadowed_image = (shadowed_image * 255.0).byte() # Convert back to [0, 255] range
136. output_image_path = os.path.join(OUTPUT_SHADOW_IMAGES_DIR, shadow_free_filename)
137. transforms.ToPILImage()(shadowed_image).save(output_image_path)
139. # 9. Save the intensity-adjusted shadow mask (inverted, for visualization)
140. # Invert the shadow tensor: shadowed area = 0, non-shadowed area = 1
141. adjusted_mask = 1.0 - shadow_tensor
142. adjusted_mask = (adjusted_mask * 255.0).byte()
143. adjusted_mask_path = os.path.join(OUTPUT_ADJUSTED_MASKS_DIR, shadow_free_filename)
144. transforms.ToPILImage()(adjusted_mask).save(adjusted_mask_path)
146. # ============================================================================
147. # PROGRAM ENTRY POINT
148. # ============================================================================
150. if __name__ == "__main__":
151. # Print program information
152. print("=" * 80)
153. print("SSD-DIS Dataset Generation Tool")
154. print("Function: Synthesize shadowed document images from shadow-free images and shadow mask library")
155. print("=" * 80)
157. # Execute the synthesis process
158. synthesize_shadow_dataset()

batch_synthesize.py at commit 1fee5d2, no license · at the source