Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -53,24 +53,20 @@ st.set_page_config(
|
|
| 53 |
}
|
| 54 |
)
|
| 55 |
|
| 56 |
-
#
|
| 57 |
-
st.session_state.setdefault('
|
| 58 |
-
st.session_state.setdefault('
|
| 59 |
-
st.session_state.setdefault('
|
| 60 |
-
st.session_state.setdefault('
|
| 61 |
-
st.session_state.setdefault('
|
| 62 |
-
st.session_state.setdefault('
|
| 63 |
-
st.session_state.setdefault('
|
| 64 |
-
st.session_state.setdefault('
|
| 65 |
-
st.session_state.setdefault('
|
| 66 |
-
st.session_state.setdefault('
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
#
|
| 70 |
-
if 'asset_gallery_container' not in st.session_state:
|
| 71 |
-
st.session_state['asset_gallery_container'] = st.sidebar.empty()
|
| 72 |
-
|
| 73 |
-
@dataclass # ModelConfig: A blueprint for model configurations.
|
| 74 |
class ModelConfig:
|
| 75 |
name: str
|
| 76 |
base_model: str
|
|
@@ -78,48 +74,43 @@ class ModelConfig:
|
|
| 78 |
domain: Optional[str] = None
|
| 79 |
model_type: str = "causal_lm"
|
| 80 |
@property
|
| 81 |
-
def model_path(self):
|
| 82 |
-
return f"models/{self.name}"
|
| 83 |
|
| 84 |
-
@dataclass # DiffusionConfig: Where diffusion magic takes shape
|
| 85 |
class DiffusionConfig:
|
| 86 |
name: str
|
| 87 |
base_model: str
|
| 88 |
size: str
|
| 89 |
domain: Optional[str] = None
|
| 90 |
@property
|
| 91 |
-
def model_path(self):
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
|
| 95 |
-
|
| 96 |
-
self.
|
| 97 |
-
self.
|
| 98 |
-
self.
|
| 99 |
-
self.jokes = [
|
| 100 |
"Why did the AI go to therapy? Too many layers to unpack! ๐",
|
| 101 |
"Training complete! Time for a binary coffee break. โ",
|
| 102 |
"I told my neural network a joke; it couldn't stop dropping bits! ๐ค",
|
| 103 |
"I asked the AI for a pun, and it said, 'I'm punning on parallel processing!' ๐",
|
| 104 |
"Debugging my code is like a stand-up routineโalways a series of exceptions! ๐"
|
| 105 |
]
|
| 106 |
-
def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
|
| 107 |
-
with st.spinner(f"Loading {model_path}... โณ"):
|
| 108 |
self.model = AutoModelForCausalLM.from_pretrained(model_path)
|
| 109 |
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
|
| 110 |
-
if self.tokenizer.pad_token is None:
|
| 111 |
-
|
| 112 |
-
if
|
| 113 |
-
|
| 114 |
-
self.model.to("cuda" if torch.cuda.is_available() else "cpu")
|
| 115 |
-
st.success(f"Model loaded! ๐ {random.choice(self.jokes)}")
|
| 116 |
return self
|
| 117 |
-
def save_model(self, path: str):
|
| 118 |
-
with st.spinner("Saving model... ๐พ"):
|
| 119 |
-
os.makedirs(os.path.dirname(path), exist_ok=True)
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
st.success(f"Model saved at {path}! โ
")
|
| 123 |
|
| 124 |
class DiffusionBuilder:
|
| 125 |
def __init__(self):
|
|
@@ -140,205 +131,144 @@ class DiffusionBuilder:
|
|
| 140 |
def generate(self, prompt: str):
|
| 141 |
return self.pipeline(prompt, num_inference_steps=20).images[0]
|
| 142 |
|
| 143 |
-
def generate_filename(sequence, ext="png"):
|
| 144 |
-
return f"{sequence}_{time.strftime('%d%m%Y%H%M%S')}.{ext}"
|
| 145 |
-
|
| 146 |
def pdf_url_to_filename(url):
|
| 147 |
-
return re.sub(r'[<>:"/\\|?*]', '_', url) + ".pdf"
|
| 148 |
-
|
| 149 |
-
def
|
| 150 |
-
|
| 151 |
-
|
| 152 |
-
def
|
| 153 |
-
|
| 154 |
-
|
| 155 |
-
|
| 156 |
-
|
| 157 |
-
def get_model_files(model_type="causal_lm"):
|
| 158 |
-
return [d for d in glob.glob("models/*" if model_type == "causal_lm" else "diffusion_models/*") if os.path.isdir(d)] or ["None"]
|
| 159 |
-
|
| 160 |
-
def get_gallery_files(file_types=["png", "pdf"]):
|
| 161 |
-
return sorted(list({f for ext in file_types for f in glob.glob(f"*.{ext}")}))
|
| 162 |
-
|
| 163 |
-
def get_pdf_files():
|
| 164 |
-
return sorted(glob.glob("*.pdf"))
|
| 165 |
-
|
| 166 |
def download_pdf(url, output_path):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 167 |
try:
|
| 168 |
-
|
| 169 |
-
if
|
| 170 |
-
|
| 171 |
-
|
| 172 |
-
|
| 173 |
-
|
| 174 |
-
|
| 175 |
-
|
| 176 |
-
|
| 177 |
-
|
| 178 |
-
|
| 179 |
-
|
| 180 |
-
|
| 181 |
-
|
| 182 |
-
|
| 183 |
-
|
| 184 |
-
|
| 185 |
-
|
| 186 |
-
|
| 187 |
-
|
| 188 |
-
|
| 189 |
-
|
| 190 |
-
|
| 191 |
-
|
| 192 |
-
|
| 193 |
-
|
| 194 |
-
|
| 195 |
-
|
| 196 |
-
|
| 197 |
-
|
| 198 |
-
|
| 199 |
-
output_file = generate_filename(f"twopage_{i}", "png")
|
| 200 |
-
pix.save(output_file)
|
| 201 |
-
output_files.append(output_file)
|
| 202 |
-
elif mode == "allpages":
|
| 203 |
-
for i in range(len(doc)):
|
| 204 |
-
page = doc[i]
|
| 205 |
-
pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
|
| 206 |
-
output_file = generate_filename(f"page_{i}", "png")
|
| 207 |
-
pix.save(output_file)
|
| 208 |
-
output_files.append(output_file)
|
| 209 |
-
doc.close()
|
| 210 |
-
elapsed = int(time.time() - start_time)
|
| 211 |
-
status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
|
| 212 |
-
return output_files
|
| 213 |
-
except Exception as e:
|
| 214 |
-
status.error(f"Failed to process PDF: {str(e)}")
|
| 215 |
-
return []
|
| 216 |
-
|
| 217 |
-
# Async OCR: Convert images to text.
|
| 218 |
-
async def process_ocr(image, output_file):
|
| 219 |
-
start_time = time.time()
|
| 220 |
-
status = st.empty()
|
| 221 |
-
status.text("Processing GOT-OCR2_0... (0s)")
|
| 222 |
-
tokenizer = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
|
| 223 |
-
model = AutoModel.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
|
| 224 |
-
temp_file = f"temp_{int(time.time())}.png"
|
| 225 |
-
image.save(temp_file)
|
| 226 |
-
result = model.chat(tokenizer, temp_file, ocr_type='ocr')
|
| 227 |
-
os.remove(temp_file)
|
| 228 |
-
elapsed = int(time.time() - start_time)
|
| 229 |
-
status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
|
| 230 |
-
async with aiofiles.open(output_file, "w") as f:
|
| 231 |
-
await f.write(result)
|
| 232 |
-
return result
|
| 233 |
-
|
| 234 |
-
# Async Image Gen: Your image genie.
|
| 235 |
-
async def process_image_gen(prompt, output_file):
|
| 236 |
-
start_time = time.time()
|
| 237 |
-
status = st.empty()
|
| 238 |
-
status.text("Processing Image Gen... (0s)")
|
| 239 |
-
pipeline = (st.session_state['builder'].pipeline
|
| 240 |
-
if st.session_state.get('builder') and isinstance(st.session_state['builder'], DiffusionBuilder)
|
| 241 |
-
and st.session_state['builder'].pipeline
|
| 242 |
-
else StableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", torch_dtype=torch.float32).to("cpu"))
|
| 243 |
-
gen_image = pipeline(prompt, num_inference_steps=20).images[0]
|
| 244 |
-
elapsed = int(time.time() - start_time)
|
| 245 |
-
status.text(f"Image Gen completed in {elapsed}s!")
|
| 246 |
-
gen_image.save(output_file)
|
| 247 |
-
return gen_image
|
| 248 |
-
|
| 249 |
-
# GPT-Image Interpreter: Turning pixels into prose!
|
| 250 |
-
def process_image_with_prompt(image, prompt, model="gpt-4o-mini", detail="auto"):
|
| 251 |
-
buffered = BytesIO()
|
| 252 |
-
image.save(buffered, format="PNG")
|
| 253 |
-
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
| 254 |
-
messages = [{
|
| 255 |
-
"role": "user",
|
| 256 |
-
"content": [
|
| 257 |
-
{"type": "text", "text": prompt},
|
| 258 |
-
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}", "detail": detail}}
|
| 259 |
-
]
|
| 260 |
-
}]
|
| 261 |
-
try:
|
| 262 |
-
response = client.chat.completions.create(model=model, messages=messages, max_tokens=300)
|
| 263 |
-
return response.choices[0].message.content
|
| 264 |
-
except Exception as e:
|
| 265 |
-
return f"Error processing image with GPT: {str(e)}"
|
| 266 |
-
|
| 267 |
-
# GPT-Text Alchemist: Merging prompt and text.
|
| 268 |
-
def process_text_with_prompt(text, prompt, model="gpt-4o-mini"):
|
| 269 |
-
messages = [{"role": "user", "content": f"{prompt}\n\n{text}"}]
|
| 270 |
try:
|
| 271 |
-
response = client.chat.completions.create(model=model, messages=messages, max_tokens=300)
|
| 272 |
-
|
| 273 |
-
|
| 274 |
-
|
| 275 |
-
|
| 276 |
-
#
|
| 277 |
-
|
| 278 |
-
#
|
| 279 |
-
|
| 280 |
-
|
| 281 |
-
st.
|
| 282 |
-
|
| 283 |
-
#
|
| 284 |
-
|
| 285 |
-
|
| 286 |
-
|
| 287 |
-
|
| 288 |
-
|
| 289 |
-
|
| 290 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 291 |
with tab_camera:
|
| 292 |
-
st.header("Camera Snap ๐ท")
|
| 293 |
-
st.subheader("Single Capture")
|
| 294 |
-
cols = st.columns(2)
|
|
|
|
| 295 |
with cols[0]:
|
| 296 |
-
cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
|
| 297 |
if cam0_img:
|
| 298 |
-
filename = generate_filename("cam0")
|
| 299 |
-
if st.session_state['cam0_file'] and os.path.exists(st.session_state['cam0_file']):
|
| 300 |
-
|
| 301 |
-
|
| 302 |
-
|
| 303 |
-
st.session_state['
|
| 304 |
-
|
| 305 |
-
st.
|
| 306 |
-
|
| 307 |
-
|
|
|
|
| 308 |
with cols[1]:
|
| 309 |
-
cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
|
| 310 |
if cam1_img:
|
| 311 |
-
filename = generate_filename("cam1")
|
| 312 |
-
if st.session_state['cam1_file'] and os.path.exists(st.session_state['cam1_file']):
|
| 313 |
-
|
| 314 |
-
|
| 315 |
-
|
| 316 |
-
st.session_state['
|
| 317 |
-
|
| 318 |
-
st.
|
| 319 |
-
|
| 320 |
-
|
| 321 |
-
|
| 322 |
-
#
|
| 323 |
with tab_download:
|
| 324 |
-
st.header("Download PDFs ๐ฅ")
|
| 325 |
-
if st.button("Examples ๐"):
|
| 326 |
-
example_urls = [
|
| 327 |
-
|
| 328 |
-
|
| 329 |
-
|
| 330 |
-
|
| 331 |
-
"https://arxiv.org/pdf/2112.10752",
|
| 332 |
-
"https://arxiv.org/pdf/2308.11236",
|
| 333 |
-
"https://arxiv.org/pdf/1706.03762",
|
| 334 |
-
"https://arxiv.org/pdf/2006.11239",
|
| 335 |
-
"https://arxiv.org/pdf/2305.11207",
|
| 336 |
-
"https://arxiv.org/pdf/2106.09685",
|
| 337 |
-
"https://arxiv.org/pdf/2005.11401",
|
| 338 |
-
"https://arxiv.org/pdf/2106.10504"
|
| 339 |
-
]
|
| 340 |
-
st.session_state['pdf_urls'] = "\n".join(example_urls)
|
| 341 |
-
url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls', ""), height=200)
|
| 342 |
if st.button("Robo-Download ๐ค"):
|
| 343 |
urls = url_input.strip().split("\n")
|
| 344 |
progress_bar = st.progress(0)
|
|
@@ -354,7 +284,8 @@ with tab_download:
|
|
| 354 |
st.session_state['downloaded_pdfs'][url] = output_path
|
| 355 |
logger.info(f"Downloaded PDF from {url} to {output_path}")
|
| 356 |
entry = f"Downloaded PDF: {output_path}"
|
| 357 |
-
st.session_state['history']
|
|
|
|
| 358 |
st.session_state['asset_checkboxes'][output_path] = True
|
| 359 |
else:
|
| 360 |
st.error(f"Failed to nab {url} ๐ฟ")
|
|
@@ -363,320 +294,249 @@ with tab_download:
|
|
| 363 |
st.session_state['downloaded_pdfs'][url] = output_path
|
| 364 |
progress_bar.progress((idx + 1) / total_urls)
|
| 365 |
status_text.text("Robo-Download complete! ๐")
|
| 366 |
-
|
|
|
|
|
|
|
|
|
|
| 367 |
if st.button("Snapshot Selected ๐ธ"):
|
| 368 |
-
selected_pdfs = [path for path in get_gallery_files()
|
|
|
|
| 369 |
if selected_pdfs:
|
| 370 |
for pdf_path in selected_pdfs:
|
| 371 |
if not os.path.exists(pdf_path):
|
| 372 |
st.warning(f"File not found: {pdf_path}. Skipping.")
|
| 373 |
continue
|
| 374 |
-
mode_key = {"Single Page (High-Res)": "single",
|
| 375 |
-
"Two Pages (High-Res)": "twopage",
|
| 376 |
"All Pages (High-Res)": "allpages"}[mode]
|
| 377 |
snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key))
|
| 378 |
for snapshot in snapshots:
|
| 379 |
st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
|
| 380 |
st.session_state['asset_checkboxes'][snapshot] = True
|
| 381 |
-
|
| 382 |
else:
|
| 383 |
st.warning("No PDFs selected for snapshotting! Check some boxes in the sidebar.")
|
| 384 |
|
| 385 |
-
|
|
|
|
| 386 |
with tab_ocr:
|
| 387 |
-
st.header("Test OCR ๐")
|
| 388 |
-
all_files = get_gallery_files()
|
| 389 |
if all_files:
|
| 390 |
-
if st.button("OCR All Assets ๐"):
|
| 391 |
-
full_text = "# OCR Results\n\n"
|
| 392 |
for file in all_files:
|
| 393 |
-
if file.endswith('.png'):
|
| 394 |
-
|
| 395 |
-
|
| 396 |
-
|
| 397 |
-
|
| 398 |
-
|
| 399 |
-
|
| 400 |
-
|
| 401 |
-
|
| 402 |
-
|
| 403 |
-
|
| 404 |
-
|
| 405 |
-
|
| 406 |
-
with open(md_output_file, "w") as f:
|
| 407 |
-
f.write(full_text)
|
| 408 |
-
st.success(f"Full OCR saved to {md_output_file}")
|
| 409 |
-
st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True)
|
| 410 |
-
selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select")
|
| 411 |
if selected_file:
|
| 412 |
-
if selected_file.endswith('.png'):
|
| 413 |
-
|
| 414 |
-
|
| 415 |
-
|
| 416 |
-
|
| 417 |
-
|
| 418 |
-
|
| 419 |
-
|
| 420 |
-
|
| 421 |
-
|
| 422 |
-
st.session_state['processing']['ocr'] =
|
| 423 |
-
|
| 424 |
-
|
| 425 |
-
st.session_state['history'].append(entry)
|
| 426 |
-
st.text_area("OCR Result", result, height=200, key="ocr_result")
|
| 427 |
-
st.success(f"OCR output saved to {output_file}")
|
| 428 |
-
st.session_state['processing']['ocr'] = False
|
| 429 |
-
if selected_file.endswith('.pdf') and st.button("OCR All Pages ๐", key="ocr_all_pages"):
|
| 430 |
-
doc = fitz.open(selected_file)
|
| 431 |
-
full_text = f"# OCR Results for {os.path.basename(selected_file)}\n\n"
|
| 432 |
for i in range(len(doc)):
|
| 433 |
-
pix = doc[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
|
| 434 |
-
|
| 435 |
-
|
| 436 |
-
|
| 437 |
-
|
| 438 |
-
|
| 439 |
-
|
| 440 |
-
|
| 441 |
-
|
| 442 |
-
f.write(full_text)
|
| 443 |
-
st.success(f"Full OCR saved to {md_output_file}")
|
| 444 |
-
st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True)
|
| 445 |
else:
|
| 446 |
-
st.warning("No assets in gallery yet. Use Camera Snap or Download PDFs!")
|
| 447 |
|
| 448 |
-
#
|
| 449 |
with tab_build:
|
| 450 |
-
st.header("Build Titan ๐ฑ")
|
| 451 |
-
model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
|
| 452 |
base_model = st.selectbox(
|
| 453 |
"Select Tiny Model",
|
| 454 |
-
["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM"
|
| 455 |
else ["OFA-Sys/small-stable-diffusion-v0", "stabilityai/stable-diffusion-2-base"]
|
| 456 |
-
)
|
| 457 |
-
model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
|
| 458 |
-
domain = st.text_input("Target Domain", "general")
|
| 459 |
-
if st.button("Download Model โฌ๏ธ"):
|
| 460 |
config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(
|
| 461 |
name=model_name, base_model=base_model, size="small", domain=domain
|
| 462 |
-
)
|
| 463 |
-
builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
|
| 464 |
-
builder.load_model(base_model, config)
|
| 465 |
-
builder.
|
| 466 |
-
st.session_state['
|
| 467 |
-
|
| 468 |
-
st.session_state['
|
| 469 |
-
st.
|
| 470 |
-
|
| 471 |
-
|
| 472 |
-
st.success(f"Model downloaded and saved to {config.model_path}! ๐")
|
| 473 |
-
st.experimental_rerun()
|
| 474 |
-
|
| 475 |
-
# ----------------- TAB: Test Image Gen -----------------
|
| 476 |
with tab_imggen:
|
| 477 |
-
st.header("Test Image Gen ๐จ")
|
| 478 |
-
all_files = get_gallery_files()
|
| 479 |
if all_files:
|
| 480 |
-
selected_file = st.selectbox("Select Image or PDF", all_files, key="gen_select")
|
| 481 |
if selected_file:
|
| 482 |
-
if selected_file.endswith('.png'):
|
| 483 |
-
image = Image.open(selected_file)
|
| 484 |
else:
|
| 485 |
-
doc = fitz.open(selected_file)
|
| 486 |
-
|
| 487 |
-
|
| 488 |
-
|
| 489 |
-
st.
|
| 490 |
-
|
| 491 |
-
|
| 492 |
-
|
| 493 |
-
st.session_state['
|
| 494 |
-
result =
|
| 495 |
-
|
| 496 |
-
st.session_state['history'].append(entry)
|
| 497 |
-
st.image(result, caption="Generated Image", use_container_width=True)
|
| 498 |
-
st.success(f"Image saved to {output_file}")
|
| 499 |
-
st.session_state['processing']['gen'] = False
|
| 500 |
else:
|
| 501 |
-
st.warning("No images or PDFs in gallery yet. Use Camera Snap or Download PDFs!")
|
|
|
|
| 502 |
|
| 503 |
-
#
|
| 504 |
with tab_pdf_process:
|
| 505 |
-
st.header("PDF Process")
|
| 506 |
-
st.subheader("Upload PDFs for GPT-based text extraction")
|
| 507 |
-
gpt_models = ["gpt-4o", "gpt-4o-mini"]
|
| 508 |
-
selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="pdf_gpt_model")
|
| 509 |
-
detail_level = st.selectbox("Detail Level", ["auto", "low", "high"], key="pdf_detail_level")
|
| 510 |
-
uploaded_pdfs = st.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader")
|
| 511 |
-
view_mode = st.selectbox("View Mode", ["Single Page", "Double Page"], key="pdf_view_mode")
|
| 512 |
-
|
| 513 |
-
|
| 514 |
-
for
|
| 515 |
-
|
| 516 |
-
|
| 517 |
-
|
| 518 |
-
|
| 519 |
try:
|
| 520 |
-
doc = fitz.open(temp_pdf_path)
|
| 521 |
-
st.write(f"Processing {pdf_file.name} with {len(doc)} pages")
|
| 522 |
-
if view_mode == "Single Page":
|
| 523 |
for i, page in enumerate(doc):
|
| 524 |
-
pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
|
| 525 |
-
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
| 526 |
-
st.image(img, caption=f"{pdf_file.name} Page {i+1}")
|
| 527 |
-
gpt_text = process_image_with_prompt(
|
| 528 |
-
|
| 529 |
-
|
| 530 |
-
|
|
|
|
|
|
|
| 531 |
for i in range(0, len(pages), 2):
|
| 532 |
-
if i+1 < len(pages):
|
| 533 |
-
pix1 = pages[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
|
| 534 |
-
|
| 535 |
-
|
| 536 |
-
|
| 537 |
-
|
| 538 |
-
|
| 539 |
-
|
| 540 |
-
|
| 541 |
-
|
| 542 |
-
|
| 543 |
-
|
| 544 |
-
|
| 545 |
-
|
| 546 |
-
|
| 547 |
-
|
| 548 |
-
|
| 549 |
-
|
| 550 |
-
|
| 551 |
-
|
| 552 |
-
|
| 553 |
-
|
| 554 |
-
|
| 555 |
-
|
| 556 |
-
output_filename
|
| 557 |
-
|
| 558 |
-
|
| 559 |
-
|
| 560 |
-
|
| 561 |
-
|
| 562 |
-
# ----------------- TAB: Image Process -----------------
|
| 563 |
with tab_image_process:
|
| 564 |
-
st.header("Image Process")
|
| 565 |
-
st.subheader("Upload Images for GPT-based OCR")
|
| 566 |
-
gpt_models = ["gpt-4o", "gpt-4o-mini"]
|
| 567 |
-
selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="img_gpt_model")
|
| 568 |
-
detail_level = st.selectbox("Detail Level", ["auto", "low", "high"], key="img_detail_level")
|
| 569 |
-
prompt_img = st.text_input("Enter prompt for image processing", "Extract the electronic text from image", key="img_process_prompt")
|
| 570 |
-
uploaded_images = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="image_process_uploader")
|
| 571 |
-
if st.button("Process Uploaded Images", key="process_images"):
|
| 572 |
-
combined_text = ""
|
| 573 |
for img_file in uploaded_images:
|
| 574 |
try:
|
| 575 |
-
img = Image.open(img_file)
|
| 576 |
-
|
| 577 |
-
gpt_text
|
| 578 |
-
|
| 579 |
-
|
| 580 |
-
|
| 581 |
-
|
| 582 |
-
|
| 583 |
-
|
| 584 |
-
|
| 585 |
-
st.markdown(get_download_link(output_filename, "text/markdown", "Download Processed Image MD"), unsafe_allow_html=True)
|
| 586 |
-
|
| 587 |
-
# ----------------- TAB: MD Gallery -----------------
|
| 588 |
with tab_md_gallery:
|
| 589 |
-
st.header("MD Gallery and GPT Processing")
|
| 590 |
-
gpt_models = ["gpt-4o", "gpt-4o-mini"]
|
| 591 |
-
selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="md_gpt_model")
|
| 592 |
-
md_files = sorted(glob.glob("*.md"))
|
| 593 |
if md_files:
|
| 594 |
-
st.subheader("Individual File Processing")
|
| 595 |
-
cols = st.columns(2)
|
| 596 |
for idx, md_file in enumerate(md_files):
|
| 597 |
with cols[idx % 2]:
|
| 598 |
-
st.write(md_file)
|
| 599 |
-
if st.button(f"Process {md_file}", key=f"process_md_{md_file}"):
|
| 600 |
try:
|
| 601 |
-
with open(md_file, "r", encoding="utf-8") as f:
|
| 602 |
-
|
| 603 |
-
|
| 604 |
-
result_text
|
| 605 |
-
|
| 606 |
-
output_filename =
|
| 607 |
-
|
| 608 |
-
|
| 609 |
-
|
| 610 |
-
|
| 611 |
-
|
| 612 |
-
st.
|
| 613 |
-
st.
|
| 614 |
-
|
| 615 |
-
|
| 616 |
-
selected_md[md_file] = st.checkbox(md_file, key=f"checkbox_md_{md_file}")
|
| 617 |
-
batch_prompt = st.text_input("Enter batch processing prompt", "Summarize this into markdown outline with emojis and number the topics 1..12", key="batch_prompt")
|
| 618 |
-
if st.button("Process Selected MD Files", key="process_batch_md"):
|
| 619 |
-
combined_content = ""
|
| 620 |
for md_file, selected in selected_md.items():
|
| 621 |
if selected:
|
| 622 |
try:
|
| 623 |
-
with open(md_file, "r", encoding="utf-8") as f:
|
| 624 |
-
|
| 625 |
-
except Exception as e:
|
| 626 |
-
st.error(f"Error reading {md_file}: {str(e)}")
|
| 627 |
if combined_content:
|
| 628 |
-
result_text = process_text_with_prompt(combined_content, batch_prompt, model=selected_gpt_model)
|
| 629 |
-
st.markdown(result_text)
|
| 630 |
-
output_filename = generate_filename("batch_processed_md", "md")
|
| 631 |
-
with open(output_filename, "w", encoding="utf-8") as f:
|
| 632 |
-
|
| 633 |
-
st.
|
| 634 |
-
st.markdown(get_download_link(output_filename, "text/markdown", "Download Batch Processed MD"), unsafe_allow_html=True)
|
| 635 |
else:
|
| 636 |
-
st.warning("No MD files selected.")
|
| 637 |
else:
|
| 638 |
-
st.warning("No MD files found.")
|
| 639 |
-
|
| 640 |
-
# ----------------- FINAL SIDEBAR UPDATE -----------------
|
| 641 |
-
# Update the asset gallery once (using its container).
|
| 642 |
-
def update_gallery():
|
| 643 |
-
container = st.session_state['asset_gallery_container']
|
| 644 |
-
container.empty() # Clear previous gallery content.
|
| 645 |
-
all_files = get_gallery_files()
|
| 646 |
-
if all_files:
|
| 647 |
-
container.markdown("### Asset Gallery ๐ธ๐")
|
| 648 |
-
cols = container.columns(2)
|
| 649 |
-
for idx, file in enumerate(all_files[:st.session_state['gallery_size']]):
|
| 650 |
-
with cols[idx % 2]:
|
| 651 |
-
st.session_state['unique_counter'] += 1
|
| 652 |
-
unique_id = st.session_state['unique_counter']
|
| 653 |
-
if file.endswith('.png'):
|
| 654 |
-
st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
|
| 655 |
-
else:
|
| 656 |
-
doc = fitz.open(file)
|
| 657 |
-
pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
|
| 658 |
-
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
| 659 |
-
st.image(img, caption=os.path.basename(file), use_container_width=True)
|
| 660 |
-
doc.close()
|
| 661 |
-
checkbox_key = f"asset_{file}_{unique_id}"
|
| 662 |
-
st.session_state['asset_checkboxes'][file] = st.checkbox("Use for SFT/Input", value=st.session_state['asset_checkboxes'].get(file, False), key=checkbox_key)
|
| 663 |
-
mime_type = "image/png" if file.endswith('.png') else "application/pdf"
|
| 664 |
-
st.markdown(get_download_link(file, mime_type, "Snag It! ๐ฅ"), unsafe_allow_html=True)
|
| 665 |
-
if st.button("Zap It! ๐๏ธ", key=f"delete_{file}_{unique_id}"):
|
| 666 |
-
os.remove(file)
|
| 667 |
-
st.session_state['asset_checkboxes'].pop(file, None)
|
| 668 |
-
st.success(f"Asset {os.path.basename(file)} vaporized! ๐จ")
|
| 669 |
-
st.experimental_rerun()
|
| 670 |
-
|
| 671 |
-
# Call the gallery update once after all tabs have been processed.
|
| 672 |
-
update_gallery()
|
| 673 |
-
|
| 674 |
-
# Finally, update the Action Logs and History in the sidebar.
|
| 675 |
-
st.sidebar.subheader("Action Logs ๐")
|
| 676 |
-
for record in log_records:
|
| 677 |
-
st.sidebar.write(f"{record.asctime} - {record.levelname} - {record.message}")
|
| 678 |
-
|
| 679 |
-
st.sidebar.subheader("History ๐")
|
| 680 |
-
for entry in st.session_state.get("history", []):
|
| 681 |
-
if entry is not None:
|
| 682 |
-
st.sidebar.write(entry)
|
|
|
|
| 53 |
}
|
| 54 |
)
|
| 55 |
|
| 56 |
+
st.session_state.setdefault('history', []) # ๐ฑ History: starting fresh if empty!
|
| 57 |
+
st.session_state.setdefault('builder', None) # ๐ ๏ธ Builder: set up the builder if it's missing!
|
| 58 |
+
st.session_state.setdefault('model_loaded', False) # ๐ฆ Model Loaded: mark as not loaded by default!
|
| 59 |
+
st.session_state.setdefault('processing', {}) # โณ Processing: initialize processing state as an empty dict!
|
| 60 |
+
st.session_state.setdefault('asset_checkboxes', {}) # โ
Asset Checkboxes: default to an empty dictionary!
|
| 61 |
+
st.session_state.setdefault('downloaded_pdfs', {}) # ๐ Downloaded PDFs: start with no PDFs downloaded!
|
| 62 |
+
st.session_state.setdefault('unique_counter', 0) # ๐ข Unique Counter: initialize the counter to zero!
|
| 63 |
+
st.session_state.setdefault('selected_model_type', "Causal LM") # ๐ง Selected Model Type: default to "Causal LM"!
|
| 64 |
+
st.session_state.setdefault('selected_model', "None") # ๐ค Selected Model: set to "None" if not already set!
|
| 65 |
+
st.session_state.setdefault('cam0_file', None) # ๐ธ Cam0 File: no file loaded by default!
|
| 66 |
+
st.session_state.setdefault('cam1_file', None) # ๐ธ Cam1 File: no file loaded by default!
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
@dataclass # ๐จ ModelConfig: A blueprint for model configurations!
|
|
|
|
|
|
|
|
|
|
|
|
|
| 70 |
class ModelConfig:
|
| 71 |
name: str
|
| 72 |
base_model: str
|
|
|
|
| 74 |
domain: Optional[str] = None
|
| 75 |
model_type: str = "causal_lm"
|
| 76 |
@property
|
| 77 |
+
def model_path(self): return f"models/{self.name}" # ๐ Model Path: Home base for brilliance!
|
|
|
|
| 78 |
|
| 79 |
+
@dataclass # ๐จ DiffusionConfig: Where diffusion magic takes shape!
|
| 80 |
class DiffusionConfig:
|
| 81 |
name: str
|
| 82 |
base_model: str
|
| 83 |
size: str
|
| 84 |
domain: Optional[str] = None
|
| 85 |
@property
|
| 86 |
+
def model_path(self): return f"diffusion_models/{self.name}" # ๐ Diffusion Path: Let the diffusion begin!
|
| 87 |
+
|
| 88 |
+
class ModelBuilder: # ๐ง ModelBuilder: Crafting AI wonders with wit!
|
| 89 |
+
def __init__(self): # ๐ Initialize: Setting up the AI factory!
|
| 90 |
+
self.config = None # No config yetโwaiting for genius!
|
| 91 |
+
self.model = None # Model not built until the magic happens!
|
| 92 |
+
self.tokenizer = None # Tokenizer: Ready to speak in AI!
|
| 93 |
+
self.jokes = [ # ๐คฃ Jokes to keep the circuits laughing!
|
|
|
|
| 94 |
"Why did the AI go to therapy? Too many layers to unpack! ๐",
|
| 95 |
"Training complete! Time for a binary coffee break. โ",
|
| 96 |
"I told my neural network a joke; it couldn't stop dropping bits! ๐ค",
|
| 97 |
"I asked the AI for a pun, and it said, 'I'm punning on parallel processing!' ๐",
|
| 98 |
"Debugging my code is like a stand-up routineโalways a series of exceptions! ๐"
|
| 99 |
]
|
| 100 |
+
def load_model(self, model_path: str, config: Optional[ModelConfig] = None): # ๐ load_model: Booting up genius!
|
| 101 |
+
with st.spinner(f"Loading {model_path}... โณ"): # โณ Spinner: Genius loading...
|
| 102 |
self.model = AutoModelForCausalLM.from_pretrained(model_path)
|
| 103 |
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
|
| 104 |
+
if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token # ๐ง Fix pad token if missing!
|
| 105 |
+
if config: self.config = config # ๐ ๏ธ Config loadedโsetting the stage!
|
| 106 |
+
self.model.to("cuda" if torch.cuda.is_available() else "cpu") # ๐ป Deploying the model to its device!
|
| 107 |
+
st.success(f"Model loaded! ๐ {random.choice(self.jokes)}") # ๐ Success: Model is now in orbit!
|
|
|
|
|
|
|
| 108 |
return self
|
| 109 |
+
def save_model(self, path: str): # ๐พ save_model: Securing your masterpiece!
|
| 110 |
+
with st.spinner("Saving model... ๐พ"): # โณ Spinner: Saving brilliance...
|
| 111 |
+
os.makedirs(os.path.dirname(path), exist_ok=True); self.model.save_pretrained(path); self.tokenizer.save_pretrained(path) # ๐ Directory magic: Creating and saving!
|
| 112 |
+
st.success(f"Model saved at {path}! โ
") # โ
Success: Your model is safely stored!
|
| 113 |
+
|
|
|
|
| 114 |
|
| 115 |
class DiffusionBuilder:
|
| 116 |
def __init__(self):
|
|
|
|
| 131 |
def generate(self, prompt: str):
|
| 132 |
return self.pipeline(prompt, num_inference_steps=20).images[0]
|
| 133 |
|
| 134 |
+
def generate_filename(sequence, ext="png"): return f"{sequence}_{time.strftime('%d%m%Y%H%M%S')}.{ext}" # โณ Generate filename with timestamp magic!
|
|
|
|
|
|
|
| 135 |
def pdf_url_to_filename(url):
|
| 136 |
+
return re.sub(r'[<>:"/\\|?*]', '_', url) + ".pdf" # ๐ Convert URL to a safe PDF filename โ no hackers allowed!
|
| 137 |
+
def get_download_link(file_path, mime_type="application/pdf", label="Download"): return f'<a href="data:{mime_type};base64,{base64.b64encode(open(file_path, "rb").read()).decode()}" download="{os.path.basename(file_path)}">{label}</a>' # ๐ Create a download link โ click it like it's hot!
|
| 138 |
+
def zip_directory(directory_path, zip_path):
|
| 139 |
+
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: [zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path))) for root, _, files in os.walk(directory_path) for file in files] # ๐ Zip directory: Packing files faster than Santa on Christmas Eve!
|
| 140 |
+
def get_model_files(model_type="causal_lm"): return [d for d in glob.glob("models/*" if model_type == "causal_lm" else "diffusion_models/*") if os.path.isdir(d)] or ["None"] # ๐ Get model files: Hunting directories like a pro!
|
| 141 |
+
def get_gallery_files(file_types=["png", "pdf"]): return sorted(list({f for ext in file_types for f in glob.glob(f"*.{ext}")})) # ๐ผ๏ธ Get gallery files: Finding art in a digital haystack!
|
| 142 |
+
def get_pdf_files(): return sorted(glob.glob("*.pdf")) # ๐ Get PDF files: Sorted and served โ no paper cuts here!
|
| 143 |
+
|
| 144 |
+
# ๐ฅ Download PDF: Delivering docs faster than a caffeinated courier!
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 145 |
def download_pdf(url, output_path):
|
| 146 |
+
try:
|
| 147 |
+
response = requests.get(url, stream=True, timeout=10); [open(output_path, "wb").write(chunk) for chunk in response.iter_content(chunk_size=8192)] if response.status_code == 200 else None; ret = True if response.status_code == 200 else False
|
| 148 |
+
except requests.RequestException as e:
|
| 149 |
+
logger.error(f"Failed to download {url}: {e}"); ret = False
|
| 150 |
+
return ret
|
| 151 |
+
|
| 152 |
+
# ๐ Async PDF Snapshot: Snap your PDF pages without blockingโjuggle pages like a ninja! ๐ฅท
|
| 153 |
+
async def process_pdf_snapshot(pdf_path, mode="single"):
|
| 154 |
+
start_time = time.time(); status = st.empty(); status.text(f"Processing PDF Snapshot ({mode})... (0s)")
|
| 155 |
try:
|
| 156 |
+
doc = fitz.open(pdf_path); output_files = []
|
| 157 |
+
if mode == "single": page = doc[0]; pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); output_file = generate_filename("single", "png"); pix.save(output_file); output_files.append(output_file)
|
| 158 |
+
elif mode == "twopage":
|
| 159 |
+
for i in range(min(2, len(doc))): page = doc[i]; pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); output_file = generate_filename(f"twopage_{i}", "png"); pix.save(output_file); output_files.append(output_file)
|
| 160 |
+
elif mode == "allpages":
|
| 161 |
+
for i in range(len(doc)): page = doc[i]; pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); output_file = generate_filename(f"page_{i}", "png"); pix.save(output_file); output_files.append(output_file)
|
| 162 |
+
doc.close(); elapsed = int(time.time() - start_time); status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!"); update_gallery(); return output_files
|
| 163 |
+
except Exception as e: status.error(f"Failed to process PDF: {str(e)}"); return []
|
| 164 |
+
|
| 165 |
+
# ๐ Async OCR: Convert images to text while your app keeps on groovin'โno blocking, just rocking! ๐ธ
|
| 166 |
+
async def process_ocr(image, output_file):
|
| 167 |
+
start_time = time.time(); status = st.empty(); status.text("Processing GOT-OCR2_0... (0s)")
|
| 168 |
+
tokenizer = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True); model = AutoModel.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
|
| 169 |
+
temp_file = f"temp_{int(time.time())}.png"; image.save(temp_file)
|
| 170 |
+
result = model.chat(tokenizer, temp_file, ocr_type='ocr'); os.remove(temp_file)
|
| 171 |
+
elapsed = int(time.time() - start_time); status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
|
| 172 |
+
async with aiofiles.open(output_file, "w") as f: await f.write(result)
|
| 173 |
+
update_gallery(); return result
|
| 174 |
+
|
| 175 |
+
# ๐ง Async Image Gen: Your image genieโwishing up pictures while the event loop keeps the party going! ๐
|
| 176 |
+
async def process_image_gen(prompt, output_file):
|
| 177 |
+
start_time = time.time(); status = st.empty(); status.text("Processing Image Gen... (0s)")
|
| 178 |
+
pipeline = st.session_state['builder'].pipeline if st.session_state.get('builder') and isinstance(st.session_state['builder'], DiffusionBuilder) and st.session_state['builder'].pipeline else StableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", torch_dtype=torch.float32).to("cpu")
|
| 179 |
+
gen_image = pipeline(prompt, num_inference_steps=20).images[0]; elapsed = int(time.time() - start_time)
|
| 180 |
+
status.text(f"Image Gen completed in {elapsed}s!"); gen_image.save(output_file); update_gallery(); return gen_image
|
| 181 |
+
|
| 182 |
+
# ๐ผ๏ธ GPT-Image Interpreter: Turning pixels into prose!
|
| 183 |
+
def process_image_with_prompt(image, prompt, model="gpt-4o-mini", detail="auto"):
|
| 184 |
+
buffered = BytesIO(); image.save(buffered, format="PNG") # ๐พ Save the image in-memory as PNGโno hard drives harmed!
|
| 185 |
+
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") # ๐ Encode image data in Base64 for secure, inline transmission!
|
| 186 |
+
messages = [{"role": "user", "content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}", "detail": detail}}]}] # ๐ฌ Build the GPT conversation with your prompt and image!
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 187 |
try:
|
| 188 |
+
response = client.chat.completions.create(model=model, messages=messages, max_tokens=300); return response.choices[0].message.content # ๐ค Invoke GPTโs magic and return its dazzling output!
|
| 189 |
+
except Exception as e: return f"Error processing image with GPT: {str(e)}" # โ ๏ธ OopsโGPT encountered a snag, so we catch and report the error!
|
| 190 |
+
|
| 191 |
+
# ๐ GPT-Text Alchemist: Merging your prompt and text into digital gold!
|
| 192 |
+
def process_text_with_prompt(text, prompt, model="gpt-4o-mini"):
|
| 193 |
+
messages = [{"role": "user", "content": f"{prompt}\n\n{text}"}] # ๐ ๏ธ Constructing the conversation input like a master wordsmith!
|
| 194 |
+
try:
|
| 195 |
+
response = client.chat.completions.create(model=model, messages=messages, max_tokens=300); return response.choices[0].message.content # ๐ค Summon GPTโs wisdom and return its brilliant answer!
|
| 196 |
+
except Exception as e: return f"Error processing text with GPT: {str(e)}" # โ ๏ธ Oops, GPT stumbledโcatching and reporting the error!
|
| 197 |
+
|
| 198 |
+
st.sidebar.subheader("Gallery Settings") # ๐จ Sidebar Gallery: Customize your creative space!
|
| 199 |
+
st.session_state.setdefault('gallery_size', 2) # ๐ง Setting default gallery size to 2 if it's missing!
|
| 200 |
+
st.session_state['gallery_size'] = st.sidebar.slider("Gallery Size", 1, 10, st.session_state['gallery_size'], key="gallery_size_slider") # ๐๏ธ Slide to adjust your gallery size and bring balance to your art!
|
| 201 |
+
|
| 202 |
+
# ๐ธ Gallery Updater: Making your assets dazzle and disappear faster than a magician's rabbit! ๐โจ
|
| 203 |
+
def update_gallery():
|
| 204 |
+
all_files = get_gallery_files() # ๐ Grab all gallery files like a digital treasure hunt!
|
| 205 |
+
if all_files: # โ
If assets are found, let the show begin!
|
| 206 |
+
st.sidebar.subheader("Asset Gallery ๐ธ๐"); cols = st.sidebar.columns(2) # ๐จ Set up a stylish 2-column layout in the sidebar!
|
| 207 |
+
for idx, file in enumerate(all_files[:st.session_state['gallery_size']]): # ๐ผ๏ธ Loop through your favorite files, limited by gallery size!
|
| 208 |
+
with cols[idx % 2]: # ๐ Alternate columnsโbecause balance is key (and funny)!
|
| 209 |
+
st.session_state['unique_counter'] += 1; unique_id = st.session_state['unique_counter'] # ๐ Increment your asset counterโevery asset gets its moment in the spotlight!
|
| 210 |
+
if file.endswith('.png'): st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True) # ๐ผ๏ธ Display the image like a masterpiece!
|
| 211 |
+
else: # ๐ For PDFs, we snap their first page like a paparazzo!
|
| 212 |
+
doc = fitz.open(file); pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5)); img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); st.image(img, caption=os.path.basename(file), use_container_width=True); doc.close()
|
| 213 |
+
checkbox_key = f"asset_{file}_{unique_id}" # ๐ Create a unique keyโbecause every asset deserves VIP treatment!
|
| 214 |
+
st.session_state['asset_checkboxes'][file] = st.checkbox("Use for SFT/Input", value=st.session_state['asset_checkboxes'].get(file, False), key=checkbox_key) # โ
Checkbox: Pick your asset for magic (or SFT)!
|
| 215 |
+
mime_type = "image/png" if file.endswith('.png') else "application/pdf" # ๐ Determine MIME typeโlike sorting your socks, but cooler!
|
| 216 |
+
st.markdown(get_download_link(file, mime_type, "Snag It! ๐ฅ"), unsafe_allow_html=True) # ๐ Provide a download linkโgrab your asset faster than a flash sale!
|
| 217 |
+
if st.button("Zap It! ๐๏ธ", key=f"delete_{file}_{unique_id}"): # โก "Zap It!" button: Because sometimes you just gotta make stuff disappear!
|
| 218 |
+
os.remove(file); st.session_state['asset_checkboxes'].pop(file, None); st.sidebar.success(f"Asset {os.path.basename(file)} vaporized! ๐จ"); st.rerun() # ๐ฅ Delete the file and refresh the galleryโpoof, it's gone!
|
| 219 |
+
#update_gallery() # ๐ Launch the gallery updateโlet the art party commence! (Joke: Why did the asset cross the road? To get zapped on the other side! ๐)
|
| 220 |
+
|
| 221 |
+
st.sidebar.subheader("Action Logs ๐") # ๐ Action Logs: Where our system whispers its secrets!
|
| 222 |
+
with st.sidebar: [st.write(f"{record.asctime} - {record.levelname} - {record.message}") for record in log_records] # ๐ Loop through log records and display them like diary entries!
|
| 223 |
+
|
| 224 |
+
st.sidebar.subheader("History ๐") # ๐ฐ๏ธ History: A walk down memory lane, one log at a time!
|
| 225 |
+
with st.sidebar: [st.write(entry) for entry in st.session_state['history']] # โณ Display every historic moment with style!
|
| 226 |
+
|
| 227 |
+
tabs = st.tabs(["Camera Snap ๐ท", "Download PDFs ๐ฅ", "Test OCR ๐", "Build Titan ๐ฑ", "Test Image Gen ๐จ", "PDF Process ๐", "Image Process ๐ผ๏ธ", "MD Gallery ๐"]) # ๐ญ Tabs: Navigate your AI universe like a boss!
|
| 228 |
+
(tab_camera, tab_download, tab_ocr, tab_build, tab_imggen, tab_pdf_process, tab_image_process, tab_md_gallery) = tabs # ๐ Unpack the tabs and get ready to exploreโbecause even tabs need to party!
|
| 229 |
+
|
| 230 |
with tab_camera:
|
| 231 |
+
st.header("Camera Snap ๐ท") # ๐ฅ Header: Letโs capture those Kodak moments!
|
| 232 |
+
st.subheader("Single Capture") # ๐ธ Subheader: One snap at a time, no double exposure!
|
| 233 |
+
cols = st.columns(2) # ๐งฉ Creating two columns for double-camera action!
|
| 234 |
+
|
| 235 |
with cols[0]:
|
| 236 |
+
cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0") # ๐ท Cam 0: Say cheese!
|
| 237 |
if cam0_img:
|
| 238 |
+
filename = generate_filename("cam0") # ๐ท๏ธ Filename for Cam 0 snapshot generated!
|
| 239 |
+
if st.session_state['cam0_file'] and os.path.exists(st.session_state['cam0_file']): os.remove(st.session_state['cam0_file']) # ๐๏ธ Out with the old Cam 0 snap!
|
| 240 |
+
with open(filename, "wb") as f: f.write(cam0_img.getvalue()) # ๐พ Saving Cam 0 image like a boss!
|
| 241 |
+
st.session_state['cam0_file'] = filename # ๐ Updating session state for Cam 0 file!
|
| 242 |
+
entry = f"Snapshot from Cam 0: {filename}" # ๐ History entry: Cam 0 snapshot recorded!
|
| 243 |
+
if entry not in st.session_state['history']:
|
| 244 |
+
st.session_state['history'] = [e for e in st.session_state['history'] if not e.startswith("Snapshot from Cam 0:")] + [entry] # ๐งน Cleaning and updating history!
|
| 245 |
+
st.image(Image.open(filename), caption="Camera 0", use_container_width=True) # ๐ผ๏ธ Displaying the fresh Cam 0 image!
|
| 246 |
+
logger.info(f"Saved snapshot from Camera 0: {filename}") # ๐ Logging: Cam 0 snapshot saved!
|
| 247 |
+
update_gallery() # ๐ Refreshing gallery to show the new snap!
|
| 248 |
+
|
| 249 |
with cols[1]:
|
| 250 |
+
cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1") # ๐ท Cam 1: Capture your best side!
|
| 251 |
if cam1_img:
|
| 252 |
+
filename = generate_filename("cam1") # ๐ท๏ธ Filename for Cam 1 snapshot generated!
|
| 253 |
+
if st.session_state['cam1_file'] and os.path.exists(st.session_state['cam1_file']): os.remove(st.session_state['cam1_file']) # ๐๏ธ Out with the old Cam 1 snap!
|
| 254 |
+
with open(filename, "wb") as f: f.write(cam1_img.getvalue()) # ๐พ Saving Cam 1 image like a pro!
|
| 255 |
+
st.session_state['cam1_file'] = filename # ๐ Updating session state for Cam 1 file!
|
| 256 |
+
entry = f"Snapshot from Cam 1: {filename}" # ๐ History entry: Cam 1 snapshot recorded!
|
| 257 |
+
if entry not in st.session_state['history']:
|
| 258 |
+
st.session_state['history'] = [e for e in st.session_state['history'] if not e.startswith("Snapshot from Cam 1:")] + [entry] # ๐งน Cleaning and updating history!
|
| 259 |
+
st.image(Image.open(filename), caption="Camera 1", use_container_width=True) # ๐ผ๏ธ Displaying the fresh Cam 1 image!
|
| 260 |
+
logger.info(f"Saved snapshot from Camera 1: {filename}") # ๐ Logging: Cam 1 snapshot saved!
|
| 261 |
+
update_gallery() # ๐ Refreshing gallery to show the new snap!
|
| 262 |
+
|
| 263 |
+
# === Tab: Download PDFs ===
|
| 264 |
with tab_download:
|
| 265 |
+
st.header("Download PDFs ๐ฅ") # ๐ฅ Header: Ready to snag PDFs like a digital ninja!
|
| 266 |
+
if st.button("Examples ๐"): # ๐ Button: Load up some scholarly URLs for instant fun!
|
| 267 |
+
example_urls = ["https://arxiv.org/pdf/2308.03892", "https://arxiv.org/pdf/1912.01703", "https://arxiv.org/pdf/2408.11039", "https://arxiv.org/pdf/2109.10282", "https://arxiv.org/pdf/2112.10752", "https://arxiv.org/pdf/2308.11236", "https://arxiv.org/pdf/1706.03762", "https://arxiv.org/pdf/2006.11239", "https://arxiv.org/pdf/2305.11207", "https://arxiv.org/pdf/2106.09685", "https://arxiv.org/pdf/2005.11401", "https://arxiv.org/pdf/2106.10504"]; st.session_state['pdf_urls'] = "\n".join(example_urls) # ๐ Examples loaded into session!
|
| 268 |
+
|
| 269 |
+
url_input = st.text_area("Enter PDF URLs (one per line)", value=st.session_state.get('pdf_urls', ""), height=200) # ๐ Text area: Paste your PDF URLs hereโno commas needed!
|
| 270 |
+
|
| 271 |
+
# --- Download PDFs Tab (modified section) ---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 272 |
if st.button("Robo-Download ๐ค"):
|
| 273 |
urls = url_input.strip().split("\n")
|
| 274 |
progress_bar = st.progress(0)
|
|
|
|
| 284 |
st.session_state['downloaded_pdfs'][url] = output_path
|
| 285 |
logger.info(f"Downloaded PDF from {url} to {output_path}")
|
| 286 |
entry = f"Downloaded PDF: {output_path}"
|
| 287 |
+
if entry not in st.session_state['history']:
|
| 288 |
+
st.session_state['history'].append(entry)
|
| 289 |
st.session_state['asset_checkboxes'][output_path] = True
|
| 290 |
else:
|
| 291 |
st.error(f"Failed to nab {url} ๐ฟ")
|
|
|
|
| 294 |
st.session_state['downloaded_pdfs'][url] = output_path
|
| 295 |
progress_bar.progress((idx + 1) / total_urls)
|
| 296 |
status_text.text("Robo-Download complete! ๐")
|
| 297 |
+
update_gallery()
|
| 298 |
+
|
| 299 |
+
|
| 300 |
+
mode = st.selectbox("Snapshot Mode", ["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (High-Res)"], key="download_mode") # ๐๏ธ Selectbox: Choose your snapshot resolution!
|
| 301 |
if st.button("Snapshot Selected ๐ธ"):
|
| 302 |
+
selected_pdfs = [path for path in get_gallery_files()
|
| 303 |
+
if path.endswith('.pdf') and st.session_state['asset_checkboxes'].get(path, False)]
|
| 304 |
if selected_pdfs:
|
| 305 |
for pdf_path in selected_pdfs:
|
| 306 |
if not os.path.exists(pdf_path):
|
| 307 |
st.warning(f"File not found: {pdf_path}. Skipping.")
|
| 308 |
continue
|
| 309 |
+
mode_key = {"Single Page (High-Res)": "single",
|
| 310 |
+
"Two Pages (High-Res)": "twopage",
|
| 311 |
"All Pages (High-Res)": "allpages"}[mode]
|
| 312 |
snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key))
|
| 313 |
for snapshot in snapshots:
|
| 314 |
st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
|
| 315 |
st.session_state['asset_checkboxes'][snapshot] = True
|
| 316 |
+
update_gallery()
|
| 317 |
else:
|
| 318 |
st.warning("No PDFs selected for snapshotting! Check some boxes in the sidebar.")
|
| 319 |
|
| 320 |
+
|
| 321 |
+
# === Tab: Test OCR ===
|
| 322 |
with tab_ocr:
|
| 323 |
+
st.header("Test OCR ๐") # ๐ Header: Time to turn images into textโmagic for your eyeballs!
|
| 324 |
+
all_files = get_gallery_files(); # ๐ Gathering all assets from the gallery!
|
| 325 |
if all_files:
|
| 326 |
+
if st.button("OCR All Assets ๐"): # ๐ Button: Blast OCR on every asset in one go!
|
| 327 |
+
full_text = "# OCR Results\n\n"; # ๐ Starting a full OCR report!
|
| 328 |
for file in all_files:
|
| 329 |
+
if file.endswith('.png'): image = Image.open(file) # ๐ผ๏ธ PNG? Open image directly!
|
| 330 |
+
else:
|
| 331 |
+
doc = fitz.open(file); pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); doc.close() # ๐ PDF? Grab a snapshot of the first page!
|
| 332 |
+
output_file = generate_filename(f"ocr_{os.path.basename(file)}", "txt"); # ๐พ Create a unique filename for the OCR text!
|
| 333 |
+
result = asyncio.run(process_ocr(image, output_file)); # ๐ค Run OCR asynchronouslyโnon-blocking wizardry!
|
| 334 |
+
full_text += f"## {os.path.basename(file)}\n\n{result}\n\n"; # ๐ Append the OCR result to the full report!
|
| 335 |
+
entry = f"OCR Test: {file} -> {output_file}"; # ๐ Log this OCR operation!
|
| 336 |
+
if entry not in st.session_state['history']: st.session_state['history'].append(entry) # โ
Update history if this entry is new!
|
| 337 |
+
md_output_file = f"full_ocr_{int(time.time())}.md"; # ๐ Generate a markdown filename for the full OCR report!
|
| 338 |
+
with open(md_output_file, "w") as f: f.write(full_text); # ๐พ Write the full OCR report to disk!
|
| 339 |
+
st.success(f"Full OCR saved to {md_output_file}"); # ๐ Success: Full OCR report is saved!
|
| 340 |
+
st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True) # ๐ Provide a download link for your OCR masterpiece!
|
| 341 |
+
selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select"); # ๐ Selectbox: Pick an asset for individual OCR!
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 342 |
if selected_file:
|
| 343 |
+
if selected_file.endswith('.png'): image = Image.open(selected_file) # ๐ผ๏ธ Open the selected PNG image!
|
| 344 |
+
else:
|
| 345 |
+
doc = fitz.open(selected_file); pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); doc.close() # ๐ For PDFs, extract a snapshot from the first page!
|
| 346 |
+
st.image(image, caption="Input Image", use_container_width=True); # ๐ผ๏ธ Display the selected asset for OCR review!
|
| 347 |
+
if st.button("Run OCR ๐", key="ocr_run"): # ๐ Button: Run OCR on the selected asset!
|
| 348 |
+
output_file = generate_filename("ocr_output", "txt"); st.session_state['processing']['ocr'] = True; # ๐พ Generate output filename and flag processing!
|
| 349 |
+
result = asyncio.run(process_ocr(image, output_file)); # ๐ค Execute OCR asynchronously!
|
| 350 |
+
entry = f"OCR Test: {selected_file} -> {output_file}"; # ๐ Create a log entry for this OCR run!
|
| 351 |
+
if entry not in st.session_state['history']: st.session_state['history'].append(entry); # โ
Update history if new!
|
| 352 |
+
st.text_area("OCR Result", result, height=200, key="ocr_result"); # ๐ Show the OCR result in a text area!
|
| 353 |
+
st.success(f"OCR output saved to {output_file}"); st.session_state['processing']['ocr'] = False # ๐ Success: OCR result saved and processing flag reset!
|
| 354 |
+
if selected_file.endswith('.pdf') and st.button("OCR All Pages ๐", key="ocr_all_pages"): # ๐ Button: Run OCR on every page of a PDF!
|
| 355 |
+
doc = fitz.open(selected_file); full_text = f"# OCR Results for {os.path.basename(selected_file)}\n\n"; # ๐ Start a report for multi-page PDF OCR!
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 356 |
for i in range(len(doc)):
|
| 357 |
+
pix = doc[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); # ๐ผ๏ธ Capture each page as an image!
|
| 358 |
+
output_file = generate_filename(f"ocr_page_{i}", "txt"); result = asyncio.run(process_ocr(image, output_file)); # ๐พ Generate filename and process OCR for the page!
|
| 359 |
+
full_text += f"## Page {i + 1}\n\n{result}\n\n"; # ๐ Append the page's OCR result to the report!
|
| 360 |
+
entry = f"OCR Test: {selected_file} Page {i + 1} -> {output_file}"; # ๐ Log this page's OCR operation!
|
| 361 |
+
if entry not in st.session_state['history']: st.session_state['history'].append(entry) # โ
Update history if this entry is new!
|
| 362 |
+
md_output_file = f"full_ocr_{os.path.basename(selected_file)}_{int(time.time())}.md"; # ๐ Create a markdown filename for the full multi-page OCR report!
|
| 363 |
+
with open(md_output_file, "w") as f: f.write(full_text); # ๐พ Write the full multi-page OCR report to disk!
|
| 364 |
+
st.success(f"Full OCR saved to {md_output_file}"); # ๐ Success: Multi-page OCR report is saved!
|
| 365 |
+
st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True) # ๐ Provide a download link for the multi-page OCR report!
|
|
|
|
|
|
|
|
|
|
| 366 |
else:
|
| 367 |
+
st.warning("No assets in gallery yet. Use Camera Snap or Download PDFs!") # โ ๏ธ Warning: Your gallery is emptyโcapture or download some assets first!
|
| 368 |
|
| 369 |
+
# === Tab: Build Titan ===
|
| 370 |
with tab_build:
|
| 371 |
+
st.header("Build Titan ๐ฑ") # ๐ฑ Header: Build your own Titanโtiny models, huge ambitions!
|
| 372 |
+
model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type") # ๐ Choose your model flavor!
|
| 373 |
base_model = st.selectbox(
|
| 374 |
"Select Tiny Model",
|
| 375 |
+
["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if model_type == "Causal LM"
|
| 376 |
else ["OFA-Sys/small-stable-diffusion-v0", "stabilityai/stable-diffusion-2-base"]
|
| 377 |
+
) # ๐ค Pick a tiny model based on your choice!
|
| 378 |
+
model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}") # ๐ท๏ธ Auto-generate a cool model name with a timestamp!
|
| 379 |
+
domain = st.text_input("Target Domain", "general") # ๐ฏ Specify your target domain (default: general)!
|
| 380 |
+
if st.button("Download Model โฌ๏ธ"): # โฌ๏ธ Button: Download your model and get ready to unleash the Titan!
|
| 381 |
config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(
|
| 382 |
name=model_name, base_model=base_model, size="small", domain=domain
|
| 383 |
+
) # ๐ Create model configuration on the fly!
|
| 384 |
+
builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder() # ๐ง Instantiate the builder for your model type!
|
| 385 |
+
builder.load_model(base_model, config); builder.save_model(config.model_path) # ๐ Load and save the modelโinstant Titan assembly!
|
| 386 |
+
st.session_state['builder'] = builder; st.session_state['model_loaded'] = True # โ๏ธ Update session state: model is now loaded!
|
| 387 |
+
st.session_state['selected_model_type'] = model_type; st.session_state['selected_model'] = config.model_path # ๐ Store your selection for posterity!
|
| 388 |
+
entry = f"Built {model_type} model: {model_name}" # ๐ Log the build event in history!
|
| 389 |
+
if entry not in st.session_state['history']: st.session_state['history'].append(entry)
|
| 390 |
+
st.success(f"Model downloaded and saved to {config.model_path}! ๐"); st.rerun() # ๐ Success: Titan built, now re-run to refresh the interface!
|
| 391 |
+
|
| 392 |
+
# === Tab: Test Image Gen ===
|
|
|
|
|
|
|
|
|
|
|
|
|
| 393 |
with tab_imggen:
|
| 394 |
+
st.header("Test Image Gen ๐จ") # ๐จ Header: Time to get creative with AI image generation!
|
| 395 |
+
all_files = get_gallery_files() # ๐ Retrieve all gallery assets for selection.
|
| 396 |
if all_files:
|
| 397 |
+
selected_file = st.selectbox("Select Image or PDF", all_files, key="gen_select") # ๐ Select an asset to spark creativity!
|
| 398 |
if selected_file:
|
| 399 |
+
if selected_file.endswith('.png'):
|
| 400 |
+
image = Image.open(selected_file) # ๐ผ๏ธ Directly open PNG images!
|
| 401 |
else:
|
| 402 |
+
doc = fitz.open(selected_file); pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0));
|
| 403 |
+
image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); doc.close() # ๐ For PDFs, extract the first page as an image!
|
| 404 |
+
st.image(image, caption="Reference Image", use_container_width=True) # ๐ผ๏ธ Display the chosen asset as reference.
|
| 405 |
+
prompt = st.text_area("Prompt", "Generate a neon superhero version of this image", key="gen_prompt") # โ๏ธ Enter a creative prompt to transform the image!
|
| 406 |
+
if st.button("Run Image Gen ๐", key="gen_run"): # ๐ Button: Ignite the image generator!
|
| 407 |
+
output_file = generate_filename("gen_output", "png"); st.session_state['processing']['gen'] = True # ๐พ Create output filename and flag processing status.
|
| 408 |
+
result = asyncio.run(process_image_gen(prompt, output_file)) # ๐ค Run the async image generationโnon-blocking magic in action!
|
| 409 |
+
entry = f"Image Gen Test: {prompt} -> {output_file}" # ๐ Log the image generation event!
|
| 410 |
+
if entry not in st.session_state['history']: st.session_state['history'].append(entry)
|
| 411 |
+
st.image(result, caption="Generated Image", use_container_width=True) # ๐ผ๏ธ Showcase the newly generated image!
|
| 412 |
+
st.success(f"Image saved to {output_file}"); st.session_state['processing']['gen'] = False # ๐ Success: Your masterpiece is saved and processing is complete!
|
|
|
|
|
|
|
|
|
|
|
|
|
| 413 |
else:
|
| 414 |
+
st.warning("No images or PDFs in gallery yet. Use Camera Snap or Download PDFs!") # โ ๏ธ Warning: No assets availableโcapture or download some first!
|
| 415 |
+
update_gallery() # ๐ Refresh the gallery to display any updates!
|
| 416 |
|
| 417 |
+
# === Updated Tab: PDF Process ===
|
| 418 |
with tab_pdf_process:
|
| 419 |
+
st.header("PDF Process") # ๐ Header: Ready to transform your PDFs into text with GPT magic!
|
| 420 |
+
st.subheader("Upload PDFs for GPT-based text extraction") # ๐ Subheader: Upload your PDFs and let the AI do the reading!
|
| 421 |
+
gpt_models = ["gpt-4o", "gpt-4o-mini"] # ๐ค GPT Models: Pick your AI wizardโmore vision-capable models may join the party!
|
| 422 |
+
selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="pdf_gpt_model") # ๐ Select your GPT model and let it work its charm!
|
| 423 |
+
detail_level = st.selectbox("Detail Level", ["auto", "low", "high"], key="pdf_detail_level") # ๐๏ธ Detail Level: Fine-tune your extractionโs precision!
|
| 424 |
+
uploaded_pdfs = st.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True, key="pdf_process_uploader") # ๐ค Uploader: Drag & drop your PDFs for processing!
|
| 425 |
+
view_mode = st.selectbox("View Mode", ["Single Page", "Double Page"], key="pdf_view_mode") # ๐ View Mode: Choose single or double page snapshots!
|
| 426 |
+
|
| 427 |
+
if st.button("Process Uploaded PDFs", key="process_pdfs"): # โ๏ธ Button: Kick off the PDF processing extravaganza!
|
| 428 |
+
combined_text = "" # ๐ Initialize a blank slate for the GPT output!
|
| 429 |
+
for pdf_file in uploaded_pdfs: # ๐ Loop through each uploaded PDF file!
|
| 430 |
+
pdf_bytes = pdf_file.read() # ๐ฅ Read the PDF bytes into memory!
|
| 431 |
+
temp_pdf_path = f"temp_{pdf_file.name}" # ๐ท๏ธ Create a temporary filename for processing!
|
| 432 |
+
with open(temp_pdf_path, "wb") as f: f.write(pdf_bytes) # ๐พ Write the PDF to a temporary file!
|
| 433 |
try:
|
| 434 |
+
doc = fitz.open(temp_pdf_path) # ๐ Open the temporary PDF document!
|
| 435 |
+
st.write(f"Processing {pdf_file.name} with {len(doc)} pages") # ๐ Log: Display file name and page count!
|
| 436 |
+
if view_mode == "Single Page": # ๐ Single Page Mode: Process each page separately!
|
| 437 |
for i, page in enumerate(doc):
|
| 438 |
+
pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); # ๐๏ธ Create a high-res pixmap of the page!
|
| 439 |
+
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); # ๐ผ๏ธ Convert the pixmap to an image!
|
| 440 |
+
st.image(img, caption=f"{pdf_file.name} Page {i+1}"); # ๐ผ๏ธ Display the page image!
|
| 441 |
+
gpt_text = process_image_with_prompt(
|
| 442 |
+
img, "Extract the electronic text from image", model=selected_gpt_model, detail=detail_level
|
| 443 |
+
); # ๐ค Run GPT to extract text from the image!
|
| 444 |
+
combined_text += f"\n## {pdf_file.name} - Page {i+1}\n\n{gpt_text}\n"; # ๐ Append the result to the combined text!
|
| 445 |
+
else: # ๐ Double Page Mode: Process pages in pairs!
|
| 446 |
+
pages = list(doc); # ๐ข Convert document pages to a list!
|
| 447 |
for i in range(0, len(pages), 2):
|
| 448 |
+
if i+1 < len(pages): # ๐ฏ Process two pages if available!
|
| 449 |
+
pix1 = pages[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); img1 = Image.frombytes("RGB", [pix1.width, pix1.height], pix1.samples); # ๐ผ๏ธ Process first page!
|
| 450 |
+
pix2 = pages[i+1].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); img2 = Image.frombytes("RGB", [pix2.width, pix2.height], pix2.samples); # ๐ผ๏ธ Process second page!
|
| 451 |
+
total_width = img1.width + img2.width; max_height = max(img1.height, img2.height); # ๐ Calculate dimensions for the combined image!
|
| 452 |
+
combined_img = Image.new("RGB", (total_width, max_height)); # ๐ผ๏ธ Create a blank canvas for the two pages!
|
| 453 |
+
combined_img.paste(img1, (0, 0)); combined_img.paste(img2, (img1.width, 0)); # ๐จ Paste the images side by side!
|
| 454 |
+
st.image(combined_img, caption=f"{pdf_file.name} Pages {i+1}-{i+2}"); # ๐ผ๏ธ Display the combined image!
|
| 455 |
+
gpt_text = process_image_with_prompt(
|
| 456 |
+
combined_img, "Extract the electronic text from image", model=selected_gpt_model, detail=detail_level
|
| 457 |
+
); # ๐ค Extract text from the combined image!
|
| 458 |
+
combined_text += f"\n## {pdf_file.name} - Pages {i+1}-{i+2}\n\n{gpt_text}\n"; # ๐ Append the result to the combined text!
|
| 459 |
+
else: # ๐น If there's an odd page out, process it solo!
|
| 460 |
+
pix = pages[i].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)); img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples); # ๐ผ๏ธ Process the single remaining page!
|
| 461 |
+
st.image(img, caption=f"{pdf_file.name} Page {i+1}"); # ๐ผ๏ธ Display the solo page image!
|
| 462 |
+
gpt_text = process_image_with_prompt(
|
| 463 |
+
img, "Extract the electronic text from image", model=selected_gpt_model, detail=detail_level
|
| 464 |
+
); # ๐ค Run GPT extraction on the solo page!
|
| 465 |
+
combined_text += f"\n## {pdf_file.name} - Page {i+1}\n\n{gpt_text}\n"; # ๐ Append the result!
|
| 466 |
+
doc.close(); # โ
Close the PDF document to free up resources!
|
| 467 |
+
except Exception as e:
|
| 468 |
+
st.error(f"Error processing {pdf_file.name}: {str(e)}"); # โ ๏ธ Error: Report any issues during processing!
|
| 469 |
+
finally:
|
| 470 |
+
os.remove(temp_pdf_path); # ๐งน Cleanup: Remove the temporary PDF file!
|
| 471 |
+
output_filename = generate_filename("processed_pdf", "md"); # ๐ท๏ธ Generate a unique filename for the Markdown output!
|
| 472 |
+
with open(output_filename, "w", encoding="utf-8") as f: f.write(combined_text); # ๐พ Write the combined GPT text to the Markdown file!
|
| 473 |
+
st.success(f"PDF processing complete. MD file saved as {output_filename}"); # ๐ Success: Notify the user of completion!
|
| 474 |
+
st.markdown(get_download_link(output_filename, "text/markdown", "Download Processed PDF MD"), unsafe_allow_html=True); # ๐ Provide a download link for your processed file!
|
| 475 |
+
|
| 476 |
+
# === Updated Tab: Image Process ===
|
|
|
|
|
|
|
| 477 |
with tab_image_process:
|
| 478 |
+
st.header("Image Process") # ๐ผ๏ธ Header: Transform images into text with GPT magic!
|
| 479 |
+
st.subheader("Upload Images for GPT-based OCR") # ๐ Subheader: Let your images speak for themselves!
|
| 480 |
+
gpt_models = ["gpt-4o", "gpt-4o-mini"] # ๐ค GPT Models: Choose your image wizard!
|
| 481 |
+
selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="img_gpt_model") # ๐ Pick your GPT model for image processing!
|
| 482 |
+
detail_level = st.selectbox("Detail Level", ["auto", "low", "high"], key="img_detail_level") # ๐๏ธ Detail Level: Set your extraction precision!
|
| 483 |
+
prompt_img = st.text_input("Enter prompt for image processing", "Extract the electronic text from image", key="img_process_prompt") # โ๏ธ Prompt: Tell GPT what to extract!
|
| 484 |
+
uploaded_images = st.file_uploader("Upload image files", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="image_process_uploader") # ๐ค Uploader: Drag & drop your images here!
|
| 485 |
+
if st.button("Process Uploaded Images", key="process_images"): # ๐ Button: Fire up the image processing!
|
| 486 |
+
combined_text = "" # ๐ Initialize combined text output!
|
| 487 |
for img_file in uploaded_images:
|
| 488 |
try:
|
| 489 |
+
img = Image.open(img_file); st.image(img, caption=img_file.name) # ๐ธ Display each uploaded image!
|
| 490 |
+
gpt_text = process_image_with_prompt(img, prompt_img, model=selected_gpt_model, detail=detail_level) # ๐ค Process image with GPT magic!
|
| 491 |
+
combined_text += f"\n## {img_file.name}\n\n{gpt_text}\n" # ๐ Append GPT output with file header!
|
| 492 |
+
except Exception as e: st.error(f"Error processing image {img_file.name}: {str(e)}") # โ ๏ธ Oops: Report errors if any!
|
| 493 |
+
output_filename = generate_filename("processed_image", "md") # ๐พ Generate a unique filename for the Markdown output!
|
| 494 |
+
with open(output_filename, "w", encoding="utf-8") as f: f.write(combined_text) # ๐ Save the combined GPT output!
|
| 495 |
+
st.success(f"Image processing complete. MD file saved as {output_filename}") # ๐ Success: Notify the user!
|
| 496 |
+
st.markdown(get_download_link(output_filename, "text/markdown", "Download Processed Image MD"), unsafe_allow_html=True) # ๐ Provide a download link!
|
| 497 |
+
|
| 498 |
+
# === Updated Tab: MD Gallery ===
|
|
|
|
|
|
|
|
|
|
| 499 |
with tab_md_gallery:
|
| 500 |
+
st.header("MD Gallery and GPT Processing") # ๐ Header: Where markdown meets GPT wizardry!
|
| 501 |
+
gpt_models = ["gpt-4o", "gpt-4o-mini"] # ๐ค GPT Models: Pick your processing partner!
|
| 502 |
+
selected_gpt_model = st.selectbox("Select GPT Model", gpt_models, key="md_gpt_model") # ๐ Select a GPT model for MD processing!
|
| 503 |
+
md_files = sorted(glob.glob("*.md")) # ๐ Gather all Markdown files in the directory!
|
| 504 |
if md_files:
|
| 505 |
+
st.subheader("Individual File Processing") # ๐ Subheader: Process files one at a time!
|
| 506 |
+
cols = st.columns(2) # ๐งฉ Set up two columns for a balanced view!
|
| 507 |
for idx, md_file in enumerate(md_files):
|
| 508 |
with cols[idx % 2]:
|
| 509 |
+
st.write(md_file) # ๐ Show the filename!
|
| 510 |
+
if st.button(f"Process {md_file}", key=f"process_md_{md_file}"): # ๐ Button: Process this file!
|
| 511 |
try:
|
| 512 |
+
with open(md_file, "r", encoding="utf-8") as f: content = f.read() # ๐ Read file content!
|
| 513 |
+
prompt_md = "Summarize this into markdown outline with emojis and number the topics 1..12" # โ๏ธ Prompt: Summarize with style!
|
| 514 |
+
result_text = process_text_with_prompt(content, prompt_md, model=selected_gpt_model) # ๐ค Let GPT work its magic!
|
| 515 |
+
st.markdown(result_text) # ๐จ Display the GPT output!
|
| 516 |
+
output_filename = generate_filename(f"processed_{os.path.splitext(md_file)[0]}", "md") # ๐พ Create a unique output filename!
|
| 517 |
+
with open(output_filename, "w", encoding="utf-8") as f: f.write(result_text) # ๐ Save the processed content!
|
| 518 |
+
st.markdown(get_download_link(output_filename, "text/markdown", f"Download {output_filename}"), unsafe_allow_html=True) # ๐ Provide a download link!
|
| 519 |
+
except Exception as e: st.error(f"Error processing {md_file}: {str(e)}") # โ ๏ธ Report errors if processing fails!
|
| 520 |
+
st.subheader("Batch Processing") # ๐ Subheader: Combine and process multiple files at once!
|
| 521 |
+
st.write("Select MD files to combine and process:") # ๐ Instruction: Choose files for batch processing!
|
| 522 |
+
selected_md = {} # ๐๏ธ Initialize selection dictionary!
|
| 523 |
+
for md_file in md_files: selected_md[md_file] = st.checkbox(md_file, key=f"checkbox_md_{md_file}") # โ
Create checkboxes for each file!
|
| 524 |
+
batch_prompt = st.text_input("Enter batch processing prompt", "Summarize this into markdown outline with emojis and number the topics 1..12", key="batch_prompt") # โ๏ธ Batch prompt: Set your summarization style!
|
| 525 |
+
if st.button("Process Selected MD Files", key="process_batch_md"): # ๐ Button: Process the selected files!
|
| 526 |
+
combined_content = "" # ๐ Initialize combined content string!
|
|
|
|
|
|
|
|
|
|
|
|
|
| 527 |
for md_file, selected in selected_md.items():
|
| 528 |
if selected:
|
| 529 |
try:
|
| 530 |
+
with open(md_file, "r", encoding="utf-8") as f: combined_content += f"\n## {md_file}\n" + f.read() + "\n" # ๐ Append each selected file's content!
|
| 531 |
+
except Exception as e: st.error(f"Error reading {md_file}: {str(e)}") # โ ๏ธ Report errors if file reading fails!
|
|
|
|
|
|
|
| 532 |
if combined_content:
|
| 533 |
+
result_text = process_text_with_prompt(combined_content, batch_prompt, model=selected_gpt_model) # ๐ค Process the batch with GPT!
|
| 534 |
+
st.markdown(result_text) # ๐จ Display the combined GPT output!
|
| 535 |
+
output_filename = generate_filename("batch_processed_md", "md") # ๐พ Generate a unique filename for the batch output!
|
| 536 |
+
with open(output_filename, "w", encoding="utf-8") as f: f.write(result_text) # ๐ Save the batch processed text!
|
| 537 |
+
st.success(f"Batch processing complete. MD file saved as {output_filename}") # ๐ Notify success!
|
| 538 |
+
st.markdown(get_download_link(output_filename, "text/markdown", "Download Batch Processed MD"), unsafe_allow_html=True) # ๐ Provide a download link!
|
|
|
|
| 539 |
else:
|
| 540 |
+
st.warning("No MD files selected.") # โ ๏ธ Warning: No files were chosen for batch processing!
|
| 541 |
else:
|
| 542 |
+
st.warning("No MD files found.") # โ ๏ธ Warning: Your gallery is emptyโno markdown files available!
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|