From ace6f641b603e54835d4c9f003e675d8b4e1e7eb Mon Sep 17 00:00:00 2001
From: Cyberes <cyberes@evulid.cc>
Date: Sat, 11 Nov 2023 11:16:03 -0700
Subject: [PATCH] more conversion handling

---
 pkg/image.py  |  8 ++++++++
 pkg/thread.py | 18 +++++++++++-------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/pkg/image.py b/pkg/image.py
index 2c4bdba..4538502 100644
--- a/pkg/image.py
+++ b/pkg/image.py
@@ -3,6 +3,7 @@ from pathlib import Path
 
 import PIL
 from PIL import Image
+from tqdm import tqdm
 
 
 def random_file_width(base_path: Path):
@@ -25,3 +26,10 @@ def is_png(file_path):
     except PIL.UnidentifiedImageError as e:
         # tqdm.write(str(e))
         return False, None
+
+
+def convert(file_path):
+    img = Image.open(file_path)
+    if img.format != 'PNG':
+        img.save(file_path, format='PNG')
+        tqdm.write(f'Converted {file_path} from {img.format} to PNG')
diff --git a/pkg/thread.py b/pkg/thread.py
index 016c342..9970a7a 100644
--- a/pkg/thread.py
+++ b/pkg/thread.py
@@ -6,7 +6,7 @@ from PIL import Image
 from tqdm import tqdm
 
 from pkg.proxies import PROXIES
-from .image import is_png
+from .image import convert, is_png
 
 
 def del_path(p: Path):
@@ -23,11 +23,17 @@ def download_tile(task):
         output_path: Path = output / f"{row}_{col}.png"
         if output_path.exists():
             valid_png_file, image_type = is_png(output_path)
-            if not valid_png_file:
+            if convert_to_png and image_type != 'PNG':
+                # The image was read sucessfully by PIL but it's in the wrong format.
+                convert(output_path)
+                if not is_png(output_path):
+                    tqdm.write(f'PNG conversion for {output_path} failed. Deleting and retrying...')
+                    corrupted_image = True
+            elif not valid_png_file:
                 # We will re-download the image. Don't need to delete it, just overwrite it.
                 # del_path(output_path)
                 corrupted_image = True
-                tqdm.write(f'Cannot identify image file: "{output_path}" (is {image_type}), deleting and retrying...')
+                tqdm.write(f'Bad image file: "{output_path}" (is format: {image_type}), deleting and retrying...')
             else:
                 return row, col, 'exist'
         tile_url = f"{base_url}/{row}/{col}".replace('//', '/').replace(':/', '://')
@@ -40,10 +46,8 @@ def download_tile(task):
                 f.write(response.content)
 
             if convert_to_png:
-                img = Image.open(output_path)
-                if img.format != 'PNG':
-                    img.save(output_path, format='PNG')
-                    tqdm.write(f'Converted {output_path} from {img.format} to PNG')
+                convert(output_path)
+                corrupted_image = True  # force a re-check
 
             # Recheck the PNG if it was corrupted.
             if corrupted_image and not is_png(output_path):