flaxdiff 0.1.21__py3-none-any.whl → 0.1.22__py3-none-any.whl

flaxdiff/data/online_loader.py

@@ -21,6 +21,7 @@ import urllib
 
 import PIL.Image
 import cv2
+import traceback 
 
 USER_AGENT = get_datasets_user_agent()
 
@@ -43,7 +44,27 @@ def fetch_single_image(image_url, timeout=None, retries=0):
     return image
 
 
-def default_image_processor(image, image_shape, interpolation=cv2.INTER_CUBIC):
+def default_image_processor(
+    image, image_shape, 
+    min_image_shape=(128, 128),
+    upscale_interpolation=cv2.INTER_CUBIC,
+    downscale_interpolation=cv2.INTER_AREA,
+):
+    image = np.array(image)
+    original_height, original_width = image.shape[:2]
+    # check if the image is too small
+    if min(original_height, original_width) < min(min_image_shape):
+        return None, original_height, original_width
+    # check if wrong aspect ratio
+    if max(original_height, original_width) / min(original_height, original_width) > 2.4:
+        return None, original_height, original_width
+    # check if the variance is too low
+    if np.std(image) < 1e-5:
+        return None, original_height, original_width
+    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    downscale = max(original_width, original_height) > max(image_shape)
+    interpolation = downscale_interpolation if downscale else upscale_interpolation
+
     image = A.longest_max_size(image, max(
         image_shape), interpolation=interpolation)
     image = A.pad(
@@ -53,7 +74,7 @@ def default_image_processor(image, image_shape, interpolation=cv2.INTER_CUBIC):
         border_mode=cv2.BORDER_CONSTANT,
         value=[255, 255, 255],
     )
-    return image
+    return image, original_height, original_width
 
 
 def map_sample(
@@ -72,23 +93,13 @@ def map_sample(
         if image is None:
             return
 
-        image = np.array(image)
-        original_height, original_width = image.shape[:2]
-        # check if the image is too small
-        if min(original_height, original_width) < min(min_image_shape):
-            return
-        # check if wrong aspect ratio
-        if max(original_height, original_width) / min(original_height, original_width) > 2.4:
-            return
-        # check if the variance is too low
-        if np.std(image) < 1e-4:
+        image, original_height, original_width = image_processor(
+            image, image_shape, min_image_shape=min_image_shape,
+            upscale_interpolation=upscale_interpolation,
+            downscale_interpolation=downscale_interpolation,)
+        
+        if image is None:
             return
-        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        downscale = max(original_width, original_height) > max(image_shape)
-        interpolation = downscale_interpolation if downscale else upscale_interpolation
-
-        image = image_processor(
-            image, image_shape, interpolation=interpolation)
         
         data_queue.put({
             "url": url,
@@ -98,7 +109,8 @@ def map_sample(
             "original_width": original_width,
         })
     except Exception as e:
-        print(f"Error processing {url}", e)
+        print(f"Error maping sample {url}", e)
+        traceback.print_exc() 
         # error_queue.put_nowait({
         #     "url": url,
         #     "caption": caption,
@@ -122,7 +134,8 @@ def map_batch(
         with ThreadPoolExecutor(max_workers=num_threads) as executor:
             executor.map(map_sample_fn, batch["url"], batch['caption'])
     except Exception as e:
-        print(f"Error processing batch", e)
+        print(f"Error maping batch", e)
+        traceback.print_exc() 
         # error_queue.put_nowait({
         #     "batch": batch,
         #     "error": str(e)

{flaxdiff-0.1.21.dist-info → flaxdiff-0.1.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flaxdiff
-Version: 0.1.21
+Version: 0.1.22
 Summary: A versatile and easy to understand Diffusion library
 Author: Ashish Kumar Singh
 Author-email: ashishkmr472@gmail.com

{flaxdiff-0.1.21.dist-info → flaxdiff-0.1.22.dist-info}/RECORD

@@ -1,7 +1,7 @@
 flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flaxdiff/utils.py,sha256=B0GcHlzlVYDNEIdh2v5qmP4u0neIT-FqexNohuyuCvg,2452
 flaxdiff/data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
-flaxdiff/data/online_loader.py,sha256=w6gi1tAzWr4gtPt7onpStzOxp7Kdo_2q8Ro4Yi7OT4w,10549
+flaxdiff/data/online_loader.py,sha256=LIK_O1C3yDPvvAEOWvsJrVeBopVqjg2IOMTbiSIvH6M,11025
 flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
 flaxdiff/models/attention.py,sha256=ZbDGIb5Q6FRqJ6qRY660cqw4WvF9IwCnhEuYdTpLPdM,13023
 flaxdiff/models/common.py,sha256=fd-Fl0VCNEBjijHNwGBqYL5VvXe9u0347h25czNTmRw,10780
@@ -34,7 +34,7 @@ flaxdiff/trainer/__init__.py,sha256=T-vUVq4zHcMK6kpCsG4Gu8vn71q6lZD-lg-Ul7yKfEk,
 flaxdiff/trainer/autoencoder_trainer.py,sha256=al7AsZ7yeDMEiDD-gbcXf0ADq_xfk1VMxvg24GfA-XQ,7008
 flaxdiff/trainer/diffusion_trainer.py,sha256=wKkg63DWZjx2MoM3VQNCDIr40rWN8fUGxH9jWWxfZao,9373
 flaxdiff/trainer/simple_trainer.py,sha256=Z77zRS5viJpd2Mpl6sonJk5WcnEWi2Cd4gl4u5tIX2M,18206
-flaxdiff-0.1.21.dist-info/METADATA,sha256=k1s_EIWBL0y4oCxXxr3QIi7LQbR47_jyDFfjVbURSMY,22083
-flaxdiff-0.1.21.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-flaxdiff-0.1.21.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.21.dist-info/RECORD,,
+flaxdiff-0.1.22.dist-info/METADATA,sha256=Sv8OtwO7oEcPUc7GfytcqLNC6GP8ZFA7BZ4-X5QqUj8,22083
+flaxdiff-0.1.22.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+flaxdiff-0.1.22.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.22.dist-info/RECORD,,