mung-bean/sceneweaver
0
1import spacy2from fastcoref import FCoref3from typing import List4import re, os5from googletrans import Translator6from dotenv import load_dotenv7import requests8from fastapi import HTTPException9 10load_dotenv()11 12HF_API_KEY = os.getenv("HF_API_KEY")13 14 15nlp = spacy.load("en_core_web_lg")16coref_model = FCoref()17CAPITALIZED_PRONOUNS = {18 "He",19 "She",20 "His",21 "Her",22 "They",23 "Their",24 "It",25 "Its",26 "You",27 "Your",28 "I",29 "We",30 "Our",31}32 33 34def clean_caption(line: str) -> str:35 # Remove leading numbers, dashes, bullets, bold text36 line = re.sub(r"^\s*\d+[\.\-)]?\s*", "", line) # "1. " or "2-" or "3)"37 line = re.sub(r"^\*\*(.*?)\*\*\s*[-–—]?\s*", "", line) # "**Title** —"38 return line.strip()39 40 41def get_script_captions(script_text: str):42 API_URL = "https://router.huggingface.co/novita/v3/openai/chat/completions"43 headers = {44 "Authorization": f"Bearer {HF_API_KEY}",45 "Content-Type": "application/json",46 }47 48 messages = [49 {50 "role": "system",51 "content": (52 "You convert story or movie scripts into detailed, visually rich image generation captions. "53 "Each caption should describe a visually distinct scene as if it were to be illustrated or rendered, "54 "but do not number them or include titles — just full, descriptive sentences."55 "add the appropriate camera shot or angle each sentence."56 "these are the 8 shots you can use, close up shot, extreme close up shot, long shot, low angle shot, high angle shot, dutch angle, over the shoulder shot, medium shot."57 ),58 },59 {60 "role": "user",61 "content": f"SCRIPT:\n{script_text}\n\nReturn only image generation captions, one per line, no numbering, no headings.",62 },63 ]64 65 payload = {66 "model": "deepseek/deepseek-v3-0324",67 "messages": messages,68 "temperature": 0.7,69 }70 71 response = requests.post(API_URL, headers=headers, json=payload)72 73 if response.status_code == 200:74 json_data = response.json()75 raw_output = json_data["choices"][0]["message"]["content"]76 lines = raw_output.strip().split("\n")77 return [clean_caption(line) for line in lines if line.strip()]78 else:79 raise HTTPException(80 status_code=500,81 detail=f"DeepSeek API error: {response.status_code} - {response.text}",82 )83 84 85def is_capitalized_pronoun(span: spacy.tokens.Span, text: str) -> bool:86 """Check if the span is a single capitalized pronoun in the original text."""87 if len(span) != 1 or span[0].pos_ != "PRON":88 return False89 # Use original casing from text90 start = span.start_char91 end = span.end_char92 original_token_text = text[start:end]93 return original_token_text[0].isupper()94 95 96def get_fastcoref_clusters(doc, text):97 preds = coref_model.predict(texts=[text])98 fast_clusters = preds[0].get_clusters(as_strings=False)99 100 converted_clusters = []101 for cluster in fast_clusters:102 new_cluster = []103 for start_char, end_char in cluster:104 span = doc.char_span(start_char, end_char)105 if span is not None:106 new_cluster.append((span.start, span.end))107 if new_cluster:108 converted_clusters.append(new_cluster)109 110 return converted_clusters111 112 113def get_span_noun_indices(doc: spacy.tokens.Doc, cluster: List[List[int]]) -> List[int]:114 spans = [doc[span[0] : span[1]] for span in cluster]115 spans_pos = [[token.pos_ for token in span] for span in spans]116 return [117 i118 for i, span_pos in enumerate(spans_pos)119 if any(pos in ["NOUN", "PROPN"] for pos in span_pos)120 ]121 122 123def get_cluster_head(124 doc: spacy.tokens.Doc, cluster: List[List[int]], noun_indices: List[int]125):126 head_idx = noun_indices[0] if noun_indices else 0127 head_start, head_end = cluster[head_idx]128 head_span = doc[head_start:head_end]129 return head_span, (head_start, head_end)130 131 132def is_containing_other_spans(span: List[int], all_spans: List[List[int]]):133 return any(s != span and s[0] >= span[0] and s[1] <= span[1] for s in all_spans)134 135 136def replace_coref_span(doc, coref_span, resolved_text, mention_span):137 start, end = coref_span138 prefix = " " if start > 0 and not doc[start - 1].whitespace_ else ""139 suffix = doc[end - 1].whitespace_ if end < len(doc) else ""140 141 resolved_text[start] = prefix + mention_span.text + suffix142 for i in range(start + 1, end):143 resolved_text[i] = ""144 145 146def improved_replace_corefs(147 doc: spacy.tokens.Doc, clusters: List[List[List[int]]], text: str148):149 resolved = [token.text_with_ws for token in doc]150 all_spans = [span for cluster in clusters for span in cluster]151 152 for cluster in clusters:153 noun_indices = get_span_noun_indices(doc, cluster)154 if not noun_indices:155 continue156 157 mention_span, mention = get_cluster_head(doc, cluster, noun_indices)158 159 for coref in cluster:160 coref_span = doc[coref[0] : coref[1]]161 if (162 coref != mention163 and not is_containing_other_spans(coref, all_spans)164 and is_capitalized_pronoun(coref_span, text)165 ):166 replace_coref_span(doc, coref, resolved, mention_span)167 168 return "".join(resolved)169 170 171def detect_and_translate_to_english(text: str) -> str:172 try:173 translator = Translator()174 detected = translator.detect(text)175 if detected.lang == "tl":176 print("[Info] Detected language: Filipino (tl). Translating to English...")177 translated = translator.translate(text, src="tl", dest="en")178 return translated.text179 return text180 except Exception as e:181 print(f"[Warning] Language detection or translation failed: {e}")182 return text183 184 185def resolve_coreferences(text: str) -> str:186 doc = nlp(text)187 clusters = get_fastcoref_clusters(doc, text)188 return improved_replace_corefs(doc, clusters, text)189 190 191def remove_dialogues(text: str) -> str:192 text = re.sub(r'(["“\']).*?\1', "", text)193 text = re.sub(r"\s{2,}", " ", text)194 return text.strip()195 196 197def get_resolved_sentences(text: str) -> List[str]:198 text = detect_and_translate_to_english(text)199 resolved_text = resolve_coreferences(text)200 no_dialogue_text = remove_dialogues(resolved_text)201 resolved_doc = nlp(no_dialogue_text)202 return [sent.text.strip() for sent in resolved_doc.sents]203