Photo upload refactor

ynr/apps/moderation_queue/forms.py

@@ -11,9 +11,6 @@
 from django.core.exceptions import ValidationError
 from django.template.loader import render_to_string
 from django.urls import reverse
-from moderation_queue.helpers import (
-    convert_image_to_png,
-)
 from people.forms.forms import StrippedCharField
 from PIL import Image as PILImage
 from utils.mail import send_mail
@@ -56,18 +53,13 @@ def clean(self):
 
     def save(self, commit):
         """
-        Before saving, resize and rotate the image as needed
-        and convert the image to a PNG. This is done while the
-        image is still an InMemoryUpload object.
+        On save, start the process of normalizing the image
         """
 
-        original_image = self.instance.image
-        png_image = convert_image_to_png(original_image)
-        filename = self.instance.image.name
-        extension = filename.split(".")[-1]
-        filename = filename.replace(extension, "png")
-        self.instance.image.save(filename, png_image, save=commit)
-        return super().save(commit=commit)
+        saved: QueuedImage = super().save(commit=commit)
+        if commit:
+            saved.start_image_processing()
+        return saved
 
 
 class UploadPersonPhotoURLForm(forms.Form):

ynr/apps/moderation_queue/helpers.py

@@ -1,95 +1,128 @@
 from io import BytesIO
 from tempfile import NamedTemporaryFile
+from typing import Optional, Tuple
 
 import requests
-from candidates.models.db import ActionType, LoggedAction
-from candidates.views.version_data import get_client_ip
-from django.http import HttpResponseRedirect
-from django.shortcuts import render
-from django.urls import reverse
 from PIL import Image as PillowImage
+from PIL import ImageOps
 
-from .models import QueuedImage
-
-
-def upload_photo_response(request, person, image_form, url_form):
-    return render(
-        request,
-        "moderation_queue/photo-upload-new.html",
-        {
-            "image_form": image_form,
-            "url_form": url_form,
-            "queued_images": QueuedImage.objects.filter(
-                person=person, decision="undecided"
-            ).order_by("created"),
-            "person": person,
-        },
-    )
+# 15MB — Rekognition's S3 object size limit
+MAX_IMAGE_BYTES = 15 * 1024 * 1024
 
 
-def image_form_valid_response(request, person, image_form):
-    # Make sure that we save the user that made the upload
-    queued_image = image_form.save(commit=False)
-    queued_image.user = request.user
-    queued_image.save()
-    # Record that action:
-    LoggedAction.objects.create(
-        user=request.user,
-        action_type=ActionType.PHOTO_UPLOAD,
-        ip_address=get_client_ip(request),
-        popit_person_new_version="",
-        person=person,
-        source=image_form.cleaned_data["justification_for_use"],
-    )
-    return HttpResponseRedirect(
-        reverse("photo-upload-success", kwargs={"person_id": person.id})
-    )
+def strip_alpha(photo):
+    if photo.mode in ("RGBA", "LA") or (
+        photo.mode == "P" and "transparency" in photo.info
+    ):
+        background = PillowImage.new("RGB", photo.size, (255, 255, 255))
+        alpha_img = photo.convert("RGBA")
+        background.paste(alpha_img, mask=alpha_img.getchannel("A"))
+        alpha_img.close()
+        return background
 
+    return photo.convert("RGB")
 
-# 15MB — Rekognition's S3 object size limit
-MAX_IMAGE_BYTES = 15 * 1024 * 1024
 
+def save_png_to_bytes(photo):
+    buf = BytesIO()
+    photo.save(buf, "PNG", optimize=True)
+    size = buf.tell()
+    buf.seek(0)
+    return buf, size
 
-def convert_image_to_png(photo):
-    # Some uploaded images are CYMK, which gives you an error when
-    # you try to write them as PNG, so convert to RGBA (this is
-    # RGBA rather than RGB so that any alpha channel (transparency)
-    # is preserved).
 
+def strip_exifdata(photo):
+    photo.info.pop("exif", None)
+    photo.info.pop("icc_profile", None)
+    photo.info.pop("xmp", None)
+    photo.info.pop("XML:com.adobe.xmp", None)
+    return photo
+
+
+def check_png_size(photo) -> Tuple[Optional[BytesIO], int]:
+    """
+    Encode `photo` as PNG.
+
+    If the size is below MAX_IMAGE_BYTES then return the BytesIO buffer,
+    otherwise delete the in-memory object.
+
+    """
+    buf = BytesIO()
+    photo.save(buf, "PNG")
+    size = buf.tell()
+
+    if size <= MAX_IMAGE_BYTES:
+        buf.seek(0)
+        return buf, size
+
+    buf.close()
+    return None, size
+
+
+def convert_image_to_png(photo):
     # If the photo is not already a PillowImage object
     # coming from the form, then we need to
     # open it as a PillowImage object before
     # converting it to RGBA.
     if not isinstance(photo, PillowImage.Image):
-        photo = PillowImage.open(photo).convert("RGBA")
-    else:
-        photo = photo.convert("RGBA")
-    converted = photo.copy().convert("RGB")
-    w, h = converted.size
+        photo = PillowImage.open(photo)
+
+    photo = ImageOps.exif_transpose(photo)
+    photo = strip_alpha(photo)
+    photo = strip_exifdata(photo)
+
+    w, h = photo.size
 
     # Render at full size first; return immediately if already within the limit.
-    bytes_obj = BytesIO()
-    converted.save(bytes_obj, "PNG")
-    if bytes_obj.tell() <= MAX_IMAGE_BYTES:
-        return bytes_obj
+    # Try full-size first.
+    png_image, size = check_png_size(photo)
+    if png_image:
+        return png_image
 
     # Binary search over scale factors (0–1) to find the largest image that
     # still encodes to <= MAX_IMAGE_BYTES.
     lo, hi = 0.0, 1.0
-    best = bytes_obj  # fallback; always replaced within a couple of iterations
-    for _ in range(20):
+    best = None
+
+    for _ in range(12):
         mid = (lo + hi) / 2
-        resized = converted.resize(
-            (max(1, int(w * mid)), max(1, int(h * mid))), PillowImage.LANCZOS
+
+        resized = photo.resize(
+            (max(1, int(w * mid)), max(1, int(h * mid))),
+            PillowImage.Resampling.LANCZOS,
         )
-        buf = BytesIO()
-        resized.save(buf, "PNG")
-        if buf.tell() <= MAX_IMAGE_BYTES:
-            lo = mid  # this scale fits — search higher
-            best = buf
+
+        try:
+            png_image, size = check_png_size(resized)
+        finally:
+            resized.close()
+
+        if png_image is not None:
+            lo = mid  # it fits, try larger
+
+            if best is not None:
+                best.close()
+
+            best = png_image
         else:
-            hi = mid  # too large — search lower
-    return best
+            hi = mid  # too large, try smaller
+
+    if best is not None:
+        best.seek(0)
+        return best
+
+    # Worst case: the above has failed to find anything so we just
+    # resize the image to _something_. This is likely too lossy, but
+    # it's a failsafe to get some sort of image.
+    small = photo.thumbnail(
+        (800, 800),
+        PillowImage.Resampling.LANCZOS,
+    )
+    try:
+        buf, _ = check_png_size(small)
+        return buf
+    finally:
+        small.close()
 
 
 class ImageDownloadException(Exception):