Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| import os | |
| import uuid | |
| import json | |
| from PIL import Image | |
| from pathlib import Path | |
| import imagehash | |
| import requests | |
| from io import BytesIO | |
| from rembg import remove, new_session | |
| # Lazy global for background-removal session (small model). | |
| _bg_session = None | |
| import mediapipe as mp # still used for future extensions; not required for detection now | |
| # Paths | |
| ROOT = Path(__file__).resolve().parent | |
| OUTPUT_ROOT = ROOT.parent / "outputs" | |
| OUTPUT_ROOT.mkdir(exist_ok=True) | |
| # Quality thresholds | |
| VAR_LAP_MIN = 60 # focus | |
| BRIGHT_MIN, BRIGHT_MAX = 30, 230 | |
| FACE_CONF_MIN = 0.85 | |
| CLEAN_BG_COLOR = (128, 128, 128, 255) | |
| # OpenCV DNN face detector | |
| MODEL_DIR = ROOT / "models" | |
| PROTOTXT = MODEL_DIR / "deploy.prototxt" | |
| CAFFEMODEL = MODEL_DIR / "res10_300x300_ssd_iter_140000.caffemodel" | |
| if not PROTOTXT.exists() or not CAFFEMODEL.exists(): | |
| raise FileNotFoundError("Face detector weights not found. Expected deploy.prototxt and res10_300x300_ssd_iter_140000.caffemodel in preprocess/models/") | |
| FACE_NET = cv2.dnn.readNetFromCaffe(str(PROTOTXT), str(CAFFEMODEL)) | |
| def pil_to_cv(img: Image.Image): | |
| return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) | |
| def cv_to_pil(img: np.ndarray): | |
| return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) | |
| def variance_of_laplacian(image): | |
| return cv2.Laplacian(image, cv2.CV_64F).var() | |
| def basic_qc(pil_img): | |
| cv_img = pil_to_cv(pil_img) | |
| gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY) | |
| focus = variance_of_laplacian(gray) | |
| mean = gray.mean() | |
| ok_focus = focus >= VAR_LAP_MIN | |
| ok_bright = BRIGHT_MIN <= mean <= BRIGHT_MAX | |
| return { | |
| "focus": round(float(focus), 2), | |
| "brightness": round(float(mean), 2), | |
| "ok_focus": ok_focus, | |
| "ok_brightness": ok_bright, | |
| } | |
| def face_detect(pil_img): | |
| cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) | |
| h, w, _ = cv_img.shape | |
| blob = cv2.dnn.blobFromImage(cv_img, 1.0, (300, 300), (104.0, 177.0, 123.0)) | |
| FACE_NET.setInput(blob) | |
| detections = FACE_NET.forward() | |
| boxes = [] | |
| for i in range(detections.shape[2]): | |
| score = detections[0, 0, i, 2] | |
| if score < FACE_CONF_MIN: | |
| continue | |
| box = detections[0, 0, i, 3:7] * np.array([w, h, w, h]) | |
| x1, y1, x2, y2 = box.astype(int) | |
| boxes.append({ | |
| "x1": int(max(0, x1)), | |
| "y1": int(max(0, y1)), | |
| "x2": int(min(w, x2)), | |
| "y2": int(min(h, y2)), | |
| "score": float(score), | |
| }) | |
| return boxes | |
| def crop_with_margin(pil_img, box, margin=0.15, square=True, target=1024): | |
| w, h = pil_img.size | |
| x1, y1, x2, y2 = box | |
| bw, bh = x2 - x1, y2 - y1 | |
| mx, my = int(bw * margin), int(bh * margin) | |
| x1 = max(0, x1 - mx) | |
| y1 = max(0, y1 - my) | |
| x2 = min(w, x2 + mx) | |
| y2 = min(h, y2 + my) | |
| if square: | |
| # make square by padding | |
| bw, bh = x2 - x1, y2 - y1 | |
| if bw > bh: | |
| diff = bw - bh | |
| y1 = max(0, y1 - diff // 2) | |
| y2 = min(h, y2 + diff - diff // 2) | |
| else: | |
| diff = bh - bw | |
| x1 = max(0, x1 - diff // 2) | |
| x2 = min(w, x2 + diff - diff // 2) | |
| cropped = pil_img.crop((x1, y1, x2, y2)).resize((target, target), Image.Resampling.LANCZOS) | |
| return cropped | |
| def get_bg_session(): | |
| global _bg_session | |
| if _bg_session is None: | |
| # u2netp is ~4 MB; default u2net is ~176 MB. | |
| _bg_session = new_session("u2netp") | |
| return _bg_session | |
| def segment_person(pil_img): | |
| # Use lightweight model; pass session explicitly to avoid arg conflicts. | |
| rgba = remove(pil_img.convert("RGB"), session=get_bg_session()) | |
| return rgba.convert("RGBA") | |
| def face_identity_crop(img: Image.Image, face_box): | |
| """ | |
| Mask everything outside the face box to neutral gray to reduce LoRA memorizing clothes/jewelry. | |
| """ | |
| x1, y1, x2, y2 = face_box | |
| w, h = img.size | |
| pad_x = int((x2 - x1) * 0.15) | |
| pad_y = int((y2 - y1) * 0.20) | |
| fx1 = max(0, x1 - pad_x) | |
| fy1 = max(0, y1 - pad_y) | |
| fx2 = min(w, x2 + pad_x) | |
| fy2 = min(h, y2 + pad_y) | |
| mask_np = np.zeros((h, w), dtype=np.uint8) | |
| mask_np[fy1:fy2, fx1:fx2] = 255 | |
| mask = Image.fromarray(mask_np) | |
| base = Image.new("RGBA", img.size, CLEAN_BG_COLOR) | |
| identity = Image.composite(img.convert("RGBA"), base, mask) | |
| return identity.convert("RGB") | |
| def florence_caption(pil_img: Image.Image): | |
| """ | |
| Optional Florence-2-Large captioning. Set FLORENCE_URL and FLORENCE_KEY env vars. | |
| """ | |
| url = os.environ.get("FLORENCE_URL") | |
| key = os.environ.get("FLORENCE_KEY") | |
| if not url or not key: | |
| return None | |
| try: | |
| buf = BytesIO() | |
| pil_img.save(buf, format="PNG") | |
| resp = requests.post( | |
| url, | |
| headers={"Authorization": f"Bearer {key}"}, | |
| files={"file": ("image.png", buf.getvalue(), "image/png")}, | |
| timeout=15, | |
| ) | |
| if resp.ok: | |
| data = resp.json() | |
| return data.get("caption") or data.get("text") | |
| except Exception: | |
| return None | |
| return None | |
| def dedupe_hashes(images): | |
| hashes = [] | |
| unique = [] | |
| for img in images: | |
| h = imagehash.phash(img) | |
| dup = False | |
| for seen in hashes: | |
| if h - seen <= 4: # Hamming distance threshold | |
| dup = True | |
| break | |
| if not dup: | |
| hashes.append(h) | |
| unique.append(img) | |
| return unique | |
| def process_images(files): | |
| if not files: | |
| return "No images uploaded", None, None | |
| imgs = [Image.open(f.name if hasattr(f, 'name') else f).convert("RGB") for f in files] | |
| # Deduplicate | |
| imgs = dedupe_hashes(imgs) | |
| processed = [] | |
| manifest = [] | |
| for idx, img in enumerate(imgs): | |
| qc = basic_qc(img) | |
| faces = face_detect(img) | |
| # prefer largest face if multiple | |
| if faces: | |
| face = sorted(faces, key=lambda b: b["score"] * ((b["x2"]-b["x1"])*(b["y2"]-b["y1"])), reverse=True)[0] | |
| crop = crop_with_margin(img, (face["x1"], face["y1"], face["x2"], face["y2"]), margin=0.15, square=True, target=1024) | |
| caption = florence_caption(crop) or "photo of sks_person, close-up, neutral lighting, straight face" | |
| else: | |
| # fallback: center crop | |
| w, h = img.size | |
| side = min(w, h) | |
| x1 = (w - side)//2 | |
| y1 = (h - side)//2 | |
| crop = img.crop((x1, y1, x1+side, y1+side)).resize((1024, 1024), Image.Resampling.LANCZOS) | |
| caption = florence_caption(crop) or "photo of sks_person" | |
| # segmentation variant | |
| cutout = segment_person(crop) | |
| # flat background | |
| bg = Image.new("RGBA", cutout.size, CLEAN_BG_COLOR) | |
| clean = Image.alpha_composite(bg, cutout).convert("RGB") | |
| # face-neutralized variant to suppress clothes/jewelry | |
| if faces: | |
| face_box = (face["x1"], face["y1"], face["x2"], face["y2"]) | |
| identity_only = face_identity_crop(clean, face_box) | |
| else: | |
| identity_only = clean | |
| uid = uuid.uuid4().hex[:8] | |
| base_dir = OUTPUT_ROOT / "preprocessed" | |
| base_dir.mkdir(parents=True, exist_ok=True) | |
| out_clean = base_dir / f"img_{idx}_{uid}_clean.png" | |
| out_orig = base_dir / f"img_{idx}_{uid}_orig.png" | |
| out_face = base_dir / f"img_{idx}_{uid}_face.png" | |
| clean.save(out_clean) | |
| crop.save(out_orig) | |
| identity_only.save(out_face) | |
| manifest.append({ | |
| "image": str(out_clean), | |
| "identity_masked": str(out_face), | |
| "caption": caption, | |
| "focus": qc["focus"], | |
| "brightness": qc["brightness"], | |
| "has_face": bool(faces), | |
| }) | |
| processed.append((crop, clean, identity_only, caption, qc)) | |
| manifest_path = OUTPUT_ROOT / "preprocessed" / "captions.jsonl" | |
| with open(manifest_path, "w") as f: | |
| for row in manifest: | |
| f.write(json.dumps(row) + "\n") | |
| # Build gallery display | |
| gallery = [] | |
| for crop, clean, identity_only, cap, qc in processed: | |
| gallery.append((crop, f"Aligned | {cap}")) | |
| gallery.append((clean, f"Clean bg | {cap} | focus {qc['focus']} | bright {qc['brightness']}")) | |
| gallery.append((identity_only, f"Face-only | {cap} | focus {qc['focus']}")) | |
| summary = f"Saved {len(processed)} images to {manifest_path.parent}. captions.jsonl created." | |
| return summary, gallery, str(manifest_path) | |
| def build_ui(): | |
| with gr.Blocks(title="LoRA Prep") as demo: | |
| gr.Markdown("# Image Preprocess for LoRA\nUpload ~8 photos (selfies + body).") | |
| with gr.Row(): | |
| inp = gr.File(file_count="multiple", file_types=["image"], label="Upload images") | |
| run_btn = gr.Button("Run Preprocess") | |
| status = gr.Markdown() | |
| gallery = gr.Gallery(label="Crops & Cleaned", columns=2, height=600) | |
| manifest_path = gr.Textbox(label="captions.jsonl path", interactive=False) | |
| run_btn.click(fn=process_images, inputs=inp, outputs=[status, gallery, manifest_path]) | |
| return demo | |
| if __name__ == "__main__": | |
| app = build_ui() | |
| # Expose locally and offer a Gradio share link for easy access if localhost is blocked. | |
| launch_info = app.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=True, | |
| show_error=True, | |
| inbrowser=False, | |
| ) | |
| try: | |
| share_url = getattr(launch_info, "share_url", None) | |
| if share_url: | |
| print(f"GRADIO_SHARE_URL: {share_url}", flush=True) | |
| except Exception: | |
| pass | |