reorder

2026-06-28 03:33:02 +02:00
parent 6ad11fc6c0
commit d522b2a267
5 changed files with 914 additions and 28 deletions
--- a/tour-comfy/edit_api.py
+++ b/tour-comfy/edit_api.py
@@ -1300,6 +1300,17 @@ def _write_all_static() -> None:
                if is_archived:
                    archived_count += 1
                
+                tags_val = p[16]
+                tags_list = []
+                if tags_val:
+                    if isinstance(tags_val, str):
+                        try:
+                            tags_list = json.loads(tags_val)
+                        except Exception:
+                            tags_list = []
+                    elif isinstance(tags_val, list):
+                        tags_list = tags_val
+
                db_images.append({
                    "filename":              p[0],
                    "name":                  p[1],
@@ -1317,6 +1328,7 @@ def _write_all_static() -> None:
                    "faceswap_source_video": p[13],
                    "archived":              is_archived,
                    "is_source":             bool(p[15]) if p[15] else False,
+                    "tags":                  tags_list,
                })
            print(f"[static] write_all: {len(db_images)} total images, {archived_count} archived")
            try:
@@ -2132,6 +2144,11 @@ def generate_video(req: VideoStitchRequest):
    if r.returncode != 0:
        raise HTTPException(500, f"ffmpeg error: {r.stderr.decode(errors='replace')[:600]}")

+    try:
+        _make_video_poster(out_path)
+    except Exception as pe:
+        print(f"[generate-video] poster extraction error: {pe}")
+
    # Register in DB so it shows up in the gallery
    person = database.get_person(req.filenames[0])
    group_id = person[1] if person and person[1] else naming.get_base_name(req.filenames[0])
@@ -2774,6 +2791,317 @@ def extract_face_endpoint(filename: str):
    return {"status": "queued", "filename": filename}


+class ImageTagsRequest(BaseModel):
+    tags: list[str]
+
+class TagActionRequest(BaseModel):
+    action: str  # "add" or "remove" or "toggle"
+    tag: str
+
+class BulkMoveRequest(BaseModel):
+    filenames: list[str]
+    target_strip: str
+
+def estimate_nsfw_21plus(filename: str) -> bool:
+    """Run WD tagger and return True if image is 21+ (NSFW/sensitive/explicit/nudity)."""
+    try:
+        output_dir = _load_output_dir()
+        fpath = os.path.join(output_dir, filename)
+        if not os.path.exists(fpath):
+            return False
+        
+        pil_img = Image.open(fpath)
+        # Run tagger with 0.2 threshold to capture sensitive/explicit content
+        tags = _run_tagger(pil_img, threshold=0.2)
+        
+        for t in tags:
+            tag_name = t["tag"].lower().replace("_", " ")
+            # Check rating category
+            if t["cat"] == 9 and tag_name in ("explicit", "questionable", "sensitive"):
+                if t["score"] >= 0.25:
+                    return True
+            # Check other explicit keywords in name
+            if any(nsfw_word in tag_name for nsfw_word in ("nude", "naked", "breasts", "pussy", "nipples", "ass", "panties", "underwear", "lingerie", "sex", "erotic")):
+                if t["score"] >= 0.3:
+                    return True
+        return False
+    except Exception as e:
+        print(f"[nsfw-estimator] Error estimating for {filename}: {e}")
+        return False
+
+@app.post("/images/{filename:path}/tags")
+def set_image_tags(filename: str, req: ImageTagsRequest):
+    person = database.get_person(filename)
+    if not person:
+        raise HTTPException(404, "Image not found")
+        
+    tags_list = list(set([str(t).upper().strip() for t in req.tags if t]))
+    
+    archived = "ARCHIVED" in tags_list
+    hidden = "HIDDEN" in tags_list
+    is_source = "SOURCE" in tags_list
+    
+    if archived:
+        if "VISIBLE" in tags_list:
+            tags_list.remove("VISIBLE")
+    else:
+        if not hidden:
+            if "VISIBLE" not in tags_list:
+                tags_list.append("VISIBLE")
+                
+    if hidden:
+        if "VISIBLE" in tags_list:
+            tags_list.remove("VISIBLE")
+    else:
+        if not archived:
+            if "VISIBLE" not in tags_list:
+                tags_list.append("VISIBLE")
+
+    conn = database.get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("""
+            UPDATE person
+            SET tags = %s,
+                archived = %s,
+                hidden = %s,
+                is_source = %s
+            WHERE filename = %s
+        """, (json.dumps(tags_list), archived, hidden, is_source, filename))
+        conn.commit()
+    finally:
+        cur.close()
+        database._put_db_connection(conn)
+        
+    _invalidate_static()
+    return {"filename": filename, "tags": tags_list}
+
+@app.post("/images/{filename:path}/tag-action")
+def tag_action_endpoint(filename: str, req: TagActionRequest):
+    person = database.get_person(filename)
+    if not person:
+        raise HTTPException(404, "not found")
+    
+    tags_val = person[2]
+    tags_list = []
+    if tags_val:
+        if isinstance(tags_val, str):
+            try:
+                tags_list = json.loads(tags_val)
+            except Exception:
+                tags_list = []
+        elif isinstance(tags_val, list):
+            tags_list = tags_val
+            
+    action = req.action.lower()
+    tag_name = req.tag.upper().strip()
+    
+    if action == "add":
+        if tag_name not in tags_list:
+            tags_list.append(tag_name)
+    elif action == "remove":
+        if tag_name in tags_list:
+            tags_list.remove(tag_name)
+    elif action == "toggle":
+        if tag_name in tags_list:
+            tags_list.remove(tag_name)
+        else:
+            tags_list.append(tag_name)
+            
+    if tag_name == "LIKE" and "LIKE" in tags_list:
+        if "DISLIKE" in tags_list:
+            tags_list.remove("DISLIKE")
+    elif tag_name == "DISLIKE" and "DISLIKE" in tags_list:
+        if "LIKE" in tags_list:
+            tags_list.remove("LIKE")
+
+    archived = "ARCHIVED" in tags_list
+    hidden = "HIDDEN" in tags_list
+    is_source = "SOURCE" in tags_list
+    
+    if archived:
+        if "VISIBLE" in tags_list:
+            tags_list.remove("VISIBLE")
+    else:
+        if not hidden:
+            if "VISIBLE" not in tags_list:
+                tags_list.append("VISIBLE")
+                
+    if hidden:
+        if "VISIBLE" in tags_list:
+            tags_list.remove("VISIBLE")
+    else:
+        if not archived:
+            if "VISIBLE" not in tags_list:
+                tags_list.append("VISIBLE")
+
+    conn = database.get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("""
+            UPDATE person
+            SET tags = %s,
+                archived = %s,
+                hidden = %s,
+                is_source = %s
+            WHERE filename = %s
+        """, (json.dumps(tags_list), archived, hidden, is_source, filename))
+        conn.commit()
+    finally:
+        cur.close()
+        database._put_db_connection(conn)
+        
+    _invalidate_static()
+    return {"filename": filename, "tags": tags_list}
+
+@app.post("/images/{filename:path}/source")
+def set_image_source(filename: str, body: dict):
+    is_source = bool(body.get("source", False))
+    person = database.get_person(filename)
+    if not person:
+        raise HTTPException(404, "not found")
+    
+    tags_val = person[2]
+    tags_list = []
+    if tags_val:
+        if isinstance(tags_val, str):
+            try:
+                tags_list = json.loads(tags_val)
+            except Exception:
+                tags_list = []
+        elif isinstance(tags_val, list):
+            tags_list = tags_val
+
+    if is_source:
+        if "SOURCE" not in tags_list:
+            tags_list.append("SOURCE")
+    else:
+        if "SOURCE" in tags_list:
+            tags_list.remove("SOURCE")
+            
+    conn = database.get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("UPDATE person SET is_source = %s, tags = %s WHERE filename = %s", (is_source, json.dumps(tags_list), filename))
+        conn.commit()
+    finally:
+        cur.close()
+        database._put_db_connection(conn)
+        
+    _invalidate_static()
+    return {"filename": filename, "is_source": is_source, "tags": tags_list}
+
+@app.post("/images/{filename:path}/estimate-21plus")
+def estimate_image_21plus(filename: str):
+    is_21plus = estimate_nsfw_21plus(filename)
+    person = database.get_person(filename)
+    if not person:
+        raise HTTPException(404, "not found")
+    
+    tags_val = person[2]
+    tags_list = []
+    if tags_val:
+        if isinstance(tags_val, str):
+            try:
+                tags_list = json.loads(tags_val)
+            except Exception:
+                tags_list = []
+        elif isinstance(tags_val, list):
+            tags_list = tags_val
+
+    if is_21plus:
+        if "21+" not in tags_list:
+            tags_list.append("21+")
+    else:
+        if "21+" in tags_list:
+            tags_list.remove("21+")
+            
+    conn = database.get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("UPDATE person SET tags = %s WHERE filename = %s", (json.dumps(tags_list), filename))
+        conn.commit()
+    finally:
+        cur.close()
+        database._put_db_connection(conn)
+        
+    _invalidate_static()
+    return {"filename": filename, "21+": is_21plus, "tags": tags_list}
+
+@app.post("/images/bulk-move")
+def bulk_move_endpoint(req: BulkMoveRequest):
+    if not req.filenames:
+        return {"status": "ok", "moved": 0}
+        
+    target = req.target_strip.upper().strip()
+    conn = database.get_db_connection()
+    cur = conn.cursor()
+    try:
+        for filename in req.filenames:
+            cur.execute("SELECT tags, archived, hidden, is_source FROM person WHERE filename = %s", (filename,))
+            row = cur.fetchone()
+            if not row:
+                continue
+            tags_val, archived, hidden, is_source = row
+            tags_list = []
+            if tags_val:
+                if isinstance(tags_val, str):
+                    try:
+                        tags_list = json.loads(tags_val)
+                    except Exception:
+                        tags_list = []
+                elif isinstance(tags_val, list):
+                    tags_list = tags_val
+            
+            if target == "VISIBLE":
+                archived = False
+                hidden = False
+                if "VISIBLE" not in tags_list:
+                    tags_list.append("VISIBLE")
+                if "ARCHIVED" in tags_list:
+                    tags_list.remove("ARCHIVED")
+                if "HIDDEN" in tags_list:
+                    tags_list.remove("HIDDEN")
+            elif target == "HIDDEN":
+                archived = False
+                hidden = True
+                if "HIDDEN" not in tags_list:
+                    tags_list.append("HIDDEN")
+                if "VISIBLE" in tags_list:
+                    tags_list.remove("VISIBLE")
+                if "ARCHIVED" in tags_list:
+                    tags_list.remove("ARCHIVED")
+            elif target == "ARCHIVED":
+                archived = True
+                hidden = False
+                if "ARCHIVED" not in tags_list:
+                    tags_list.append("ARCHIVED")
+                if "VISIBLE" in tags_list:
+                    tags_list.remove("VISIBLE")
+                if "HIDDEN" in tags_list:
+                    tags_list.remove("HIDDEN")
+            elif target == "SOURCE":
+                is_source = True
+                if "SOURCE" not in tags_list:
+                    tags_list.append("SOURCE")
+                    
+            cur.execute("""
+                UPDATE person
+                SET tags = %s,
+                    archived = %s,
+                    hidden = %s,
+                    is_source = %s
+                WHERE filename = %s
+            """, (json.dumps(tags_list), archived, hidden, is_source, filename))
+        conn.commit()
+    finally:
+        cur.close()
+        database._put_db_connection(conn)
+        
+    _invalidate_static()
+    return {"status": "ok", "moved": len(req.filenames)}
+
+
 class FaceSimilarRequest(BaseModel):
    group_id: str
    limit: int = 12