Update app.py

app.py

@@ -1,8 +1,7 @@
 import os
 import re
-import math
-from dataclasses import dataclass, asdict
-from typing import Dict, Any, Optional, Tuple
+from dataclasses import dataclass
+from typing import Dict, Any, Optional, Tuple, List
 
 import gradio as gr
 from PIL import Image, ImageStat
@@ -27,19 +26,6 @@ NEGATIVE_BASELINE = (
     "deformed hands, blurry, depth map artifacts, harsh HDR, unrealistic colors"
 )
 
-DEFAULT_SDXL_SETTINGS = {
-    "steps": 34,
-    "cfg": 5,
-    "sampler": "DPM++ SDE Karras",
-    "resolution": "1024 on long side",
-    "refiner": 0.25,
-    "hires": "1.5–2.0x upscale for micro‑detail"
-}
-
-# Lightweight captioner (free CPU). If unavailable, reverse will degrade gracefully.
-CAPTION_MODEL = "Salesforce/blip-image-captioning-base"  # CPU-friendly
-
-
 @dataclass
 class CameraSpec:
     cameraBody: str = ""
@@ -47,7 +33,6 @@ class CameraSpec:
     aperture: str = ""
     iso: Optional[int] = None
 
-
 @dataclass
 class PromptFields:
     subject: str = ""
@@ -63,17 +48,19 @@ class PromptFields:
     aspectRatio: str = "4:5"
     negatives: str = NEGATIVE_BASELINE
     model: str = "sdxl"  # "mj" | "sdxl" | "dalle"
+    # MJ
     settings_mj_s: int = 100
     settings_mj_chaos: int = 5
     settings_mj_seed: int = 42
+    # SDXL
     settings_sdxl_steps: int = 34
     settings_sdxl_cfg: int = 5
     settings_sdxl_sampler: str = "DPM++ SDE Karras"
     settings_sdxl_resolution: str = "1024x1280"
     settings_sdxl_refiner: float = 0.25
+    # DALL·E
     settings_dalle_resolution: str = "1024x1024"
 
-
 def realism_string(enabled: bool) -> str:
     if not enabled:
         return ""
@@ -83,26 +70,17 @@ def realism_string(enabled: bool) -> str:
         "subtle chromatic aberration, vignette."
     )
 
-
-def safe_join(parts):
+def safe_join(parts: List[str]) -> str:
     return " ".join([p.strip() for p in parts if p and str(p).strip()]).replace("  ", " ").strip()
 
-
 def build_universal(f: PromptFields) -> str:
-    # First sentence
-    s1 = ""
-    if f.subject:
-        s1 = f"Photo of {f.subject}"
-    else:
-        s1 = "Photo"
-
+    s1 = f"Photo of {f.subject}" if f.subject else "Photo"
     if f.environment:
         s1 += f" in/at {f.environment}"
     if f.timeWeather:
         s1 += f", {f.timeWeather}"
     s1 += "."
 
-    # Camera
     cam_bits = []
     if f.camera and f.camera.focalLengthMm:
         cam_bits.append(f"{f.camera.focalLengthMm}mm lens")
@@ -110,10 +88,7 @@ def build_universal(f: PromptFields) -> str:
         cam_bits.append(f"at {f.camera.aperture}")
     if f.camera and f.camera.iso:
         cam_bits.append(f"ISO {f.camera.iso}")
-    if cam_bits:
-        s2 = "Shot with a " + ", ".join(cam_bits) + "."
-    else:
-        s2 = ""
+    s2 = "Shot with a " + ", ".join(cam_bits) + "." if cam_bits else ""
 
     s3 = f"{f.composition}." if f.composition else ""
     s4 = f"Lighting: {f.lighting}." if f.lighting else ""
@@ -122,14 +97,11 @@ def build_universal(f: PromptFields) -> str:
     s7 = f"Color & grade: {f.colorGrade}." if f.colorGrade else ""
     s8 = realism_string(f.realismCues)
 
-    universal = safe_join([s1, s2, s3, s4, s5, s6, s7, s8])
-    return universal
-
+    return safe_join([s1, s2, s3, s4, s5, s6, s7, s8])
 
 def format_midjourney(universal: str, f: PromptFields) -> str:
     return f"{universal} --style raw --ar {f.aspectRatio} --s {f.settings_mj_s} --chaos {f.settings_mj_chaos} --seed {f.settings_mj_seed}"
 
-
 def format_sdxl(universal: str, f: PromptFields) -> Dict[str, Any]:
     return {
         "positive": universal,
@@ -144,12 +116,10 @@ def format_sdxl(universal: str, f: PromptFields) -> Dict[str, Any]:
         }
     }
 
-
 def format_dalle(universal: str, f: PromptFields) -> Dict[str, Any]:
     prose = f"A high‑resolution photograph. {universal}"
     return {"prompt": prose, "resolution": f.settings_dalle_resolution}
 
-
 def compose(
     subject, environment, timeWeather,
     cameraBody, focalLengthMm, aperture, iso,
@@ -187,26 +157,33 @@ def compose(
         settings_sdxl_refiner=float(sdxl_refiner) if str(sdxl_refiner).strip() else 0.25,
         settings_dalle_resolution=dalle_resolution or "1024x1024",
     )
-
     universal = build_universal(f)
     mj = format_midjourney(universal, f)
     sdxl = format_sdxl(universal, f)
     dalle = format_dalle(universal, f)
     return universal, mj, sdxl, dalle, (f.negatives or NEGATIVE_BASELINE)
 
-
 # ---------- Reverse prompt helpers ----------
 
+CAPTIONER = None
 def init_captioner():
     if not HAS_TRANSFORMERS:
         return None
     try:
-        return pipeline("image-to-text", model=CAPTION_MODEL)
+        return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
     except Exception:
         return None
-
 CAPTIONER = init_captioner()
 
+OBJDET = None
+def init_objdet():
+    if not HAS_TRANSFORMERS:
+        return None
+    try:
+        return pipeline("object-detection", model="facebook/detr-resnet-50")
+    except Exception:
+        return None
+OBJDET = init_objdet()
 
 def download_haarcascade() -> Optional[str]:
     if cv2 is None:
@@ -225,7 +202,6 @@ def download_haarcascade() -> Optional[str]:
     except Exception:
         return None
 
-
 def detect_faces(pil_img: Image.Image) -> int:
     if cv2 is None:
         return 0
@@ -241,28 +217,75 @@ def detect_faces(pil_img: Image.Image) -> int:
     except Exception:
         return 0
 
-
 def avg_brightness(pil_img: Image.Image) -> float:
     stat = ImageStat.Stat(pil_img.convert("L"))
     return float(stat.mean[0])
 
-
 def nearest_aspect(w: int, h: int) -> str:
     target = w / h
-    candidates = {
-        "1:1": 1.0,
-        "4:5": 0.8,
-        "5:4": 1.25,
-        "4:3": 1.333,
-        "3:2": 1.5,
-        "16:9": 1.777
-    }
-    best = min(candidates.items(), key=lambda kv: abs(kv[1] - target))[0]
-    return best
-
+    candidates = { "1:1": 1.0, "4:5": 0.8, "5:4": 1.25, "4:3": 1.333, "3:2": 1.5, "16:9": 1.777 }
+    return min(candidates.items(), key=lambda kv: abs(kv[1] - target))[0]
+
+def _article(word: str) -> str:
+    return "an" if word and word[0].lower() in "aeiou" else "a"
+
+def _label_to_phrase(label: str) -> str:
+    nice = {"tv": "television", "cell phone": "phone", "sports ball": "ball", "potted plant": "potted plant"}
+    word = nice.get(label, label)
+    return f"{_article(word)} {word}"
+
+def _centrality_score(cx, cy, W, H):
+    dx = abs(cx - W/2) / (W/2)
+    dy = abs(cy - H/2) / (H/2)
+    dist = min(1.0, (dx*dx + dy*dy) ** 0.5)
+    return 1.0 - dist
+
+def _detect_main_subject(img: Image.Image):
+    if OBJDET is None:
+        return None, []
+    try:
+        dets = OBJDET(img)
+    except Exception:
+        return None, []
+    if not dets:
+        return None, []
+
+    W, H = img.size
+    scored = []
+    for d in dets:
+        box = d.get("box", {})
+        xmin, ymin = box.get("xmin", 0), box.get("ymin", 0)
+        xmax, ymax = box.get("xmax", 0), box.get("ymax", 0)
+        w, h = max(1, xmax - xmin), max(1, ymax - ymin)
+        area = (w * h) / float(W * H)
+        cx, cy = xmin + w/2, ymin + h/2
+        central = _centrality_score(cx, cy, W, H)
+        conf = float(d.get("score", 0.0))
+        label = d.get("label", "")
+        score = conf * (0.6 * area + 0.4 * central)
+        scored.append({"label": label, "score": score})
+
+    scored.sort(key=lambda x: x["score"], reverse=True)
+    main_phrase = _label_to_phrase(scored[0]["label"])
+
+    suggestions, seen = [], set()
+    for s in scored:
+        p = _label_to_phrase(s["label"])
+        if p not in seen:
+            suggestions.append(p)
+            seen.add(p)
+        if len(suggestions) >= 5:
+            break
+    return main_phrase, suggestions
+
+def _action_from_caption(caption: str) -> str:
+    c = (caption or "").lower()
+    for key in ["running", "sprinting", "walking", "standing", "jumping", "riding", "driving", "sitting"]:
+        if key in c:
+            return key
+    return ""
 
 def extract_fields_from_image(img: Image.Image) -> Dict[str, Any]:
-    # Caption
     caption = ""
     if CAPTIONER:
         try:
@@ -271,38 +294,55 @@ def extract_fields_from_image(img: Image.Image) -> Dict[str, Any]:
                 caption = out[0].get("generated_text", "")
         except Exception:
             caption = ""
-    # Brightness -> time
+
     brightness = avg_brightness(img)
-    timeWeather = "golden hour" if 130 <= brightness <= 170 else ("daylight" if brightness > 140 else "night with ambient light")
+    if brightness > 140:
+        timeWeather = "daylight"
+    elif 100 < brightness <= 140:
+        timeWeather = "overcast daylight"
+    else:
+        timeWeather = "night with ambient light"
+
+    subject_phrase, subject_suggestions = _detect_main_subject(img)
 
-    # Faces -> portrait heuristics
     faces = detect_faces(img)
-    if faces > 0:
-        subject = "a person (no identity) " + (f"— {caption}" if caption else "")
-        composition = "eye‑level portrait, head-and-shoulders, shallow DOF, rule of thirds"
-        focal = 85
-        aperture = "f/1.8"
+    if not subject_phrase and faces > 0:
+        subject_phrase = "a person"
+
+    if not subject_phrase:
+        m = re.search(r"(a|an|the)\s+([^,.]+?)(?:\s+(on|in|at|by|with|near|amid|from)\b|[.,]|$)", (caption or "").lower())
+        subject_phrase = m.group(0).rstrip(",.") if m else ("a person" if faces > 0 else "a real-world subject")
+
+    if subject_phrase.startswith(("a person", "an person")):
+        act = _action_from_caption(caption)
+        if act and act not in subject_phrase:
+            subject_phrase = f"{subject_phrase} {act}"
+
+    if subject_phrase.startswith(("a person", "an person")):
+        focal = 35
+        aperture = "f/2.8"
         iso = 200 if "day" in timeWeather else 800
-        color_grade = "warm Portra‑like, soft contrast, high dynamic range"
-        lighting = "soft key at 45°, gentle reflector fill, subtle rim; ~5400K" if "day" in timeWeather else "soft practicals, dim ambient, ~3200–4000K"
-        micro = "skin pores, peach fuzz, flyaway hairs, natural imperfections"
-        motion = "no visible motion blur, creamy circular bokeh"
+        composition = "eye‑level, rear three‑quarter or profile, leading lines, shallow DOF"
+        lighting = "soft natural light" if "day" in timeWeather else "mixed ambient light with practicals, soft shadows"
+        micro = "skin pores, fabric textures, scuffs, dust in the air"
+        motion = "slight motion blur on limbs if running" if "running" in subject_phrase else "no visible motion blur"
+        color_grade = "neutral, true-to-life colors, gentle contrast, high micro‑contrast"
     else:
-        subject = caption or "a real-world scene"
-        composition = "eye‑level, balanced framing, leading lines, shallow DOF"
         focal = 35
         aperture = "f/2.8"
         iso = 200 if "day" in timeWeather else 800
-        color_grade = "neutral, true-to-life colors, gentle contrast, high micro-contrast"
-        lighting = "soft natural light, mild shadows" if "day" in timeWeather else "mixed ambient light with practicals, soft shadows"
+        composition = "eye‑level, balanced framing, leading lines, shallow DOF"
+        lighting = "soft natural light" if "day" in timeWeather else "mixed ambient light with practicals, soft shadows"
         micro = "texture of materials, dust, subtle scratches, specular highlights"
         motion = "slight motion blur if present, volumetric light if applicable"
+        color_grade = "neutral, true-to-life colors, gentle contrast, high micro‑contrast"
 
     w, h = img.size
     aspect = nearest_aspect(w, h)
 
-    fields = {
-        "subject": subject,
+    return {
+        "subject": subject_phrase,
+        "subjectCandidates": subject_suggestions,
         "environment": "",
         "timeWeather": timeWeather,
         "camera": {
@@ -321,15 +361,12 @@ def extract_fields_from_image(img: Image.Image) -> Dict[str, Any]:
         "negatives": NEGATIVE_BASELINE,
         "model": "sdxl"
     }
-    return fields
-
 
 def reverse_prompt(image: Image.Image):
     if image is None:
-        return gr.update(value=""), "", {"positive": "", "negative": "", "settings": {}}, {"prompt": "", "resolution": ""}, NEGATIVE_BASELINE
+        return {}, "", "", {"positive": "", "negative": "", "settings": {}}, {"prompt": "", "resolution": ""}, NEGATIVE_BASELINE, gr.update(choices=[], value=None)
 
     fields = extract_fields_from_image(image)
-    # Build objects
     f = PromptFields(
         subject=fields["subject"],
         environment=fields.get("environment", ""),
@@ -353,8 +390,10 @@ def reverse_prompt(image: Image.Image):
     mj = format_midjourney(universal, f)
     sdxl = format_sdxl(universal, f)
     dalle = format_dalle(universal, f)
-    return fields, universal, mj, sdxl, dalle, (fields.get("negatives") or NEGATIVE_BASELINE)
 
+    cands = fields.get("subjectCandidates", []) or []
+    dd = gr.update(choices=cands, value=(cands[0] if cands else None))
+    return fields, universal, mj, sdxl, dalle, (fields.get("negatives") or NEGATIVE_BASELINE), dd
 
 # ---------- Presets ----------
 
@@ -417,7 +456,6 @@ PRESETS = {
     )
 }
 
-
 def load_preset(name: str):
     f = PRESETS.get(name)
     if not f:
@@ -433,7 +471,6 @@ def load_preset(name: str):
         f.settings_dalle_resolution
     )
 
-
 # ---------- UI ----------
 
 with gr.Blocks(title=APP_TITLE) as demo:
@@ -445,7 +482,7 @@ with gr.Blocks(title=APP_TITLE) as demo:
                 preset = gr.Dropdown(choices=list(PRESETS.keys()), label="Presets")
                 load_btn = gr.Button("Load preset")
 
-                subject = gr.Textbox(label="Subject", placeholder="e.g., a 30‑year‑old person with freckles")
+                subject = gr.Textbox(label="Subject", placeholder="e.g., a person running")
                 environment = gr.Textbox(label="Environment/Setting", placeholder="e.g., sunlit loft by a large window")
                 timeWeather = gr.Textbox(label="Time & Weather", placeholder="e.g., golden hour")
 
@@ -455,8 +492,8 @@ with gr.Blocks(title=APP_TITLE) as demo:
                     aperture = gr.Textbox(label="Aperture", placeholder="e.g., f/1.8")
                     iso = gr.Textbox(label="ISO", placeholder="e.g., 200")
 
-                composition = gr.Textbox(label="Composition & Perspective", placeholder="e.g., eye‑level half‑body, rule of thirds, shallow DOF")
-                lighting = gr.Textbox(label="Lighting", placeholder="e.g., soft window key at 45°, reflector fill, subtle rim, 5400K")
+                composition = gr.Textbox(label="Composition & Perspective", placeholder="e.g., eye‑level, shallow DOF, rule of thirds")
+                lighting = gr.Textbox(label="Lighting", placeholder="e.g., soft window key at 45°, reflector fill, rim, 5400K")
                 microDetails = gr.Textbox(label="Materials & Micro‑detail", placeholder="e.g., skin pores, fabric weave, subtle scratches")
                 motionAtmosphere = gr.Textbox(label="Motion/Atmosphere", placeholder="e.g., slight motion blur, volumetric light, haze")
                 colorGrade = gr.Textbox(label="Color & Grade", placeholder="e.g., warm Portra‑like, soft contrast, high DR")
@@ -512,9 +549,10 @@ with gr.Blocks(title=APP_TITLE) as demo:
         )
 
     with gr.Tab("Reverse (Image → Prompt)"):
-        gr.Markdown("Upload an image. The app will infer fields without identifying real people, then build prompts.")
+        gr.Markdown("Upload an image. The app will infer fields without identifying real people, then build prompts. Use the detected-subject dropdown to set the main subject.")
         image_in = gr.Image(type="pil", label="Upload image")
         analyze_btn = gr.Button("Analyze & Generate")
+        subject_pick = gr.Dropdown(label="Detected subjects (pick one)", choices=[], value=None)
         fields_out = gr.JSON(label="Extracted fields (editable in Build tab if needed)")
         universal_out_r = gr.Textbox(label="Universal prompt", lines=6)
         mj_out_r = gr.Textbox(label="Midjourney prompt", lines=6)
@@ -525,21 +563,20 @@ with gr.Blocks(title=APP_TITLE) as demo:
         analyze_btn.click(
             reverse_prompt,
             inputs=[image_in],
-            outputs=[fields_out, universal_out_r, mj_out_r, sdxl_out_r, dalle_out_r, neg_out_r]
+            outputs=[fields_out, universal_out_r, mj_out_r, sdxl_out_r, dalle_out_r, neg_out_r, subject_pick]
         )
 
+        def use_picked_subject(picked):
+            return picked or ""
+
+        subject_pick.change(use_picked_subject, inputs=[subject_pick], outputs=[subject])
+
     gr.Markdown(
         "Tips\n"
-        "- For Midjourney, prepend 1–2 reference image URLs to match look/lighting; keep --style raw.\n"
+        "- For Midjourney, prepend 1–2 reference image URLs; keep --style raw.\n"
         "- For SDXL, use Refiner at 0.2–0.4 and upscale 1.5–2.0x for micro‑detail.\n"
-        "- DALL·E 3 responds best to concise, photographic prose with lens + lighting."
+        "- DALL·E 3 responds best to concise photographic prose with lens + lighting."
     )
 
 if __name__ == "__main__":
-    demo.launch()
-    gradio>=4.40.0
-pillow
-numpy
-transformers>=4.42.0
-torch>=2.3.0
-opencv-python-headless
+    demo.launch()