Update app.py

app.py

@@ -1,29 +1,75 @@
+import os
+import tempfile
+
 import gradio as gr
 import cv2
 import numpy as np
 from PIL import Image, ImageCms
-import tempfile
-import os
-
-# --- ICC PROFILE (cached once) ---
-_SRGB_ICC = ImageCms.ImageCmsProfile(ImageCms.createProfile("sRGB")).tobytes()
 
-# --- 1. THE PROCESSING ENGINE ---
+# =========================================================
+# NEW: Color profile assets + caching
+# =========================================================
+# sRGB profile is always available via Pillow
+SRGB_PROFILE = ImageCms.createProfile("sRGB")
+SRGB_ICC_BYTES = ImageCms.ImageCmsProfile(SRGB_PROFILE).tobytes()
+
+# NEW FILE YOU WILL CREATE IN YOUR SPACE:
+# profiles/AdobeRGB1998.icc
+ADOBE_ICC_PATH = os.path.join("profiles", "AdobeRGB1998.icc")
+ADOBE_PROFILE = ImageCms.getOpenProfile(ADOBE_ICC_PATH) if os.path.exists(ADOBE_ICC_PATH) else None
+
+
+def convert_and_tag(img: Image.Image, target_color_space: str):
+    """
+    NEW: Converts pixel values to the selected target profile AND returns the correct ICC bytes.
+    This is what makes sRGB vs AdobeRGB(1998) "real" for stock delivery.
+    """
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+
+    # For most AI/web images, treating the working space as sRGB is the practical default.
+    # (If you want full source-ICC detection later, we can add it.)
+    if target_color_space == "sRGB (Stock Standard)":
+        return img, SRGB_ICC_BYTES
+
+    if target_color_space == "Adobe RGB (1998)":
+        # If the profile file isn't present, we safely fall back to sRGB to avoid untagged/mismatched exports.
+        if ADOBE_PROFILE is None:
+            return img, SRGB_ICC_BYTES
+
+        # Convert from sRGB -> AdobeRGB1998
+        xform = ImageCms.buildTransformFromOpenProfiles(
+            SRGB_PROFILE,
+            ADOBE_PROFILE,
+            "RGB",
+            "RGB",
+            renderingIntent=0  # perceptual
+        )
+        out = ImageCms.applyTransform(img, xform)
+        return out, ADOBE_PROFILE.tobytes()
+
+    # Fallback safety
+    return img, SRGB_ICC_BYTES
+
+
+# =========================================================
+# 1) THE PROCESSING ENGINE
+# =========================================================
 def process_image(image, output_format, color_space, fix_lighting, add_grain):
     if image is None:
         return None, None
 
-    # Convert PIL -> numpy (RGB)
+    # Convert PIL -> numpy
     img_array = np.array(image)
 
-    # Drop alpha if present (stock-safe)
+    # Stock safety: drop alpha if present
     if img_array.ndim == 3 and img_array.shape[2] == 4:
         img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
 
-    # Convert RGB -> BGR for OpenCV ops
+    # RGB -> BGR for OpenCV
     img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
 
-    # A) LIGHTING FIX (gentle CLAHE)
+    # A) LIGHTING FIX (CLAHE) - gentle
     if fix_lighting:
         lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
         l, a, b = cv2.split(lab)
@@ -32,11 +78,12 @@ def process_image(image, output_format, color_space, fix_lighting, add_grain):
         lab_fixed = cv2.merge((l_fixed, a, b))
         img_bgr = cv2.cvtColor(lab_fixed, cv2.COLOR_LAB2BGR)
 
-    # B) TEXTURE FIX (Monochromatic Film Grain) — fixed + faster
+    # B) TEXTURE FIX (Monochromatic Film Grain) - FIXED + better
     if add_grain:
         h, w = img_bgr.shape[:2]
 
-        # Use int16 so negative noise doesn't wrap (uint8 would wrap!)
+        # IMPORTANT FIX:
+        # Use int16 noise so negative noise doesn't wrap like uint8.
         noise = np.random.normal(loc=0.0, scale=2.5, size=(h, w)).astype(np.int16)
 
         img_i16 = img_bgr.astype(np.int16)
@@ -46,38 +93,48 @@ def process_image(image, output_format, color_space, fix_lighting, add_grain):
 
         img_bgr = np.clip(img_i16, 0, 255).astype(np.uint8)
 
-    # C) Convert back to PIL (RGB)
+    # Convert back to PIL RGB
     img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
     final_pil = Image.fromarray(img_rgb).convert("RGB")
 
-    # D) Color profile tagging (keep it stock-safe)
-    save_kwargs = {}
-    if color_space == "sRGB (Stock Standard)":
-        save_kwargs["icc_profile"] = _SRGB_ICC
-    elif color_space == "Adobe RGB (1998)":
-        # Without an actual AdobeRGB1998.icc file, safest fallback is still sRGB tagging.
-        # (Prevents "untagged" rejections.)
-        save_kwargs["icc_profile"] = _SRGB_ICC
+    # C) COLOR SPACE: REAL convert + embed ICC
+    final_pil, icc_bytes = convert_and_tag(final_pil, color_space)
 
-    # E) Write a real downloadable file (this is what prevents “WebP only”)
+    # D) EXPORT: create a real file for download (prevents "webp-only" confusion)
     suffix = ".jpg" if output_format == "JPEG" else ".png"
     tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
     tmp_path = tmp.name
     tmp.close()
 
     if output_format == "JPEG":
-        final_pil.save(tmp_path, format="JPEG", quality=100, subsampling=0, optimize=True, **save_kwargs)
+        # Stock-friendly: no chroma subsampling for max quality
+        final_pil.save(
+            tmp_path,
+            format="JPEG",
+            quality=100,
+            subsampling=0,
+            optimize=True,
+            icc_profile=icc_bytes,
+        )
     else:
-        # compress_level 6 is a good speed/size balance
-        final_pil.save(tmp_path, format="PNG", compress_level=6, optimize=True, **save_kwargs)
+        # Balanced speed/size
+        final_pil.save(
+            tmp_path,
+            format="PNG",
+            compress_level=6,
+            optimize=True,
+            icc_profile=icc_bytes,
+        )
 
     # Return:
-    # - Preview image (Gradio may serve this as WebP in the UI; that’s fine)
-    # - File path for true PNG/JPEG download
+    # - Preview image (Gradio may serve preview as WebP; that's just UI)
+    # - Download file (REAL PNG/JPEG)
     return final_pil, tmp_path
 
 
-# --- 2. THE UI (Gradio) ---
+# =========================================================
+# 2) THE UI (Gradio)
+# =========================================================
 css = """
 #run-btn {background-color: #ff7c00 !important; color: white !important;}
 """
@@ -85,6 +142,12 @@ css = """
 with gr.Blocks(title="StockFix AI", css=css) as app:
     gr.Markdown("## StockFix AI: De-Plasticizer")
 
+    # NEW: We show a small note if Adobe ICC is missing
+    if ADOBE_PROFILE is None:
+        gr.Markdown(
+            "⚠️ **Adobe RGB (1998) ICC not found.** Add `profiles/AdobeRGB1998.icc` to enable true Adobe RGB exports."
+        )
+
     with gr.Row():
         with gr.Column():
             input_img = gr.Image(type="pil", label="Input AI Image")
@@ -110,6 +173,7 @@ with gr.Blocks(title="StockFix AI", css=css) as app:
 
         with gr.Column():
             output_img = gr.Image(label="Preview", type="pil")
+            # NEW: This creates a real downloadable file (PNG/JPEG)
             output_file = gr.File(label="Download (PNG/JPEG)")
 
     btn_run.click(