from flask import Flask, render_template, Response, flash, redirect, url_for, request, jsonify import cv2 import numpy as np from unstructured.partition.pdf import partition_pdf import json, base64, io, os from PIL import Image, ImageEnhance, ImageDraw from imutils.perspective import four_point_transform from dotenv import load_dotenv import pytesseract from transformers import AutoProcessor, AutoModelForImageTextToText, AutoModelForVision2Seq from langchain_community.document_loaders.image_captions import ImageCaptionLoader from werkzeug.utils import secure_filename import tempfile import torch from langchain_groq import ChatGroq from langgraph.prebuilt import create_react_agent load_dotenv() # os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY") groq_api_key = os.getenv("GROQ_API_KEY") llm = ChatGroq( model="meta-llama/llama-4-maverick-17b-128e-instruct", temperature=0, max_tokens=None, ) app = Flask(__name__) pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe" poppler_path=r"C:\poppler-23.11.0\Library\bin" count = 0 PDF_GET = r"E:\Pratham\2025\Harsh Sir\Scratch Vision\images\scratch_crab.pdf" OUTPUT_FOLDER = "OUTPUTS" DETECTED_IMAGE_FOLDER_PATH = os.path.join(OUTPUT_FOLDER,"DETECTED_IMAGE") IMAGE_FOLDER_PATH = os.path.join(OUTPUT_FOLDER, "SCANNED_IMAGE") JSON_FOLDER_PATH = os.path.join(OUTPUT_FOLDER, "EXTRACTED_JSON") for path in [OUTPUT_FOLDER, IMAGE_FOLDER_PATH, DETECTED_IMAGE_FOLDER_PATH, JSON_FOLDER_PATH]: os.makedirs(path, exist_ok=True) # Model Initialization try: smolvlm256m_processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct") # smolvlm256m_model = AutoModelForImageTextToText.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct").to("cpu") smolvlm256m_model = AutoModelForVision2Seq.from_pretrained( "HuggingFaceTB/SmolVLM-256M-Instruct", torch_dtype=torch.bfloat16 if hasattr(torch, "bfloat16") else torch.float32, _attn_implementation="eager" ).to("cpu") except Exception as e: raise RuntimeError(f"❌ Failed to load SmolVLM model: {str(e)}") # SmolVLM Image Captioning functioning def get_smolvlm_caption(image: Image.Image, prompt: str = "") -> str: try: # Ensure exactly one token if "" not in prompt: prompt = f" {prompt.strip()}" num_image_tokens = prompt.count("") if num_image_tokens != 1: raise ValueError(f"Prompt must contain exactly 1 token. Found {num_image_tokens}") inputs = smolvlm256m_processor(images=[image], text=[prompt], return_tensors="pt").to("cpu") output_ids = smolvlm256m_model.generate(**inputs, max_new_tokens=100) return smolvlm256m_processor.decode(output_ids[0], skip_special_tokens=True) except Exception as e: return f"❌ Error during caption generation: {str(e)}" # --- FUNCTION: Extract images from saved PDF --- def extract_images_from_pdf(pdf_path, output_json_path): ''' Extract images from PDF and generate structured sprite JSON ''' try: pdf_filename = os.path.splitext(os.path.basename(pdf_path))[0] # e.g., "scratch_crab" pdf_dir_path = os.path.dirname(pdf_path).replace("/", "\\") # Create subfolders extracted_image_subdir = os.path.join(DETECTED_IMAGE_FOLDER_PATH, pdf_filename) json_subdir = os.path.join(JSON_FOLDER_PATH, pdf_filename) os.makedirs(extracted_image_subdir, exist_ok=True) os.makedirs(json_subdir, exist_ok=True) # Output paths output_json_path = os.path.join(json_subdir, "extracted.json") final_json_path = os.path.join(json_subdir, "extracted_sprites.json") try: elements = partition_pdf( filename=pdf_path, strategy="hi_res", extract_image_block_types=["Image"], extract_image_block_to_payload=True, # Set to True to get base64 in output ) except Exception as e: raise RuntimeError(f"❌ Failed to extract images from PDF: {str(e)}") try: with open(output_json_path, "w") as f: json.dump([element.to_dict() for element in elements], f, indent=4) except Exception as e: raise RuntimeError(f"❌ Failed to write extracted.json: {str(e)}") try: # Display extracted images with open(output_json_path, 'r') as file: file_elements = json.load(file) except Exception as e: raise RuntimeError(f"❌ Failed to read extracted.json: {str(e)}") # Prepare manipulated sprite JSON structure manipulated_json = {} # SET A SYSTEM PROMPT system_prompt = """ You are an expert in visual scene understanding. Your Job is to analyze an image and respond acoording if asked for name give simple name by analyzing it and if ask for descrption generate a short description covering its elements. Guidelines: - Focus only the images given in Square Shape. - Don't Consider Blank areas in Image as. - Don't include generic summary or explanation outside the fields. Return only string. """ agent = create_react_agent( model = llm, tools = [], prompt = system_prompt ) # If JSON already exists, load it and find the next available Sprite number if os.path.exists(final_json_path): with open(final_json_path, "r") as existing_file: manipulated = json.load(existing_file) # Determine the next available index (e.g., Sprite 4 if 1–3 already exist) existing_keys = [int(k.replace("Sprite ", "")) for k in manipulated.keys()] start_count = max(existing_keys, default=0) + 1 else: start_count = 1 sprite_count = start_count for i,element in enumerate(file_elements): if "image_base64" in element["metadata"]: try: image_data = base64.b64decode(element["metadata"]["image_base64"]) image = Image.open(io.BytesIO(image_data)).convert("RGB") image.show(title=f"Extracted Image {i+1}") image_path = os.path.join(extracted_image_subdir, f"Sprite_{i+1}.png") image.save(image_path) with open(image_path, "rb") as image_file: image_bytes = image_file.read() img_base64 = base64.b64encode(image_bytes).decode("utf-8") # description = get_smolvlm_caption(image, prompt="Give a brief Description") # name = get_smolvlm_caption(image, prompt="give a short name/title of this Image.") def clean_caption_output(raw_output: str, prompt: str) -> str: answer = raw_output.replace(prompt, '').replace("", '').strip(" :-\n") return answer prompt_description = "Give a brief Captioning." prompt_name = "give a short name caption of this Image." content1 = [ { "type": "text", "text": f"{prompt_description}" }, { "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{img_base64}" } } ] response1 = agent.invoke({"messages": [{"role": "user", "content":content1}]}) print(response1) description = response1["messages"][-1].content content2 = [ { "type": "text", "text": f"{prompt_name}" }, { "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{img_base64}" } } ] response2 = agent.invoke({"messages": [{"role": "user", "content":content2}]}) print(response2) name = response2["messages"][-1].content #raw_description = get_smolvlm_caption(image, prompt=prompt_description) #raw_name = get_smolvlm_caption(image, prompt=prompt_name) #description = clean_caption_output(raw_description, prompt_description) #name = clean_caption_output(raw_name, prompt_name) manipulated_json[f"Sprite {sprite_count}"] = { "name": name, "base64": element["metadata"]["image_base64"], "file-path": pdf_dir_path, "description":description } sprite_count += 1 except Exception as e: print(f"⚠️ Error processing Sprite {i+1}: {str(e)}") # Save manipulated JSON with open(final_json_path, "w") as sprite_file: json.dump(manipulated_json, sprite_file, indent=4) print(f"✅ Manipulated sprite JSON saved: {final_json_path}") return final_json_path, manipulated_json except Exception as e: raise RuntimeError(f"❌ Error in extract_images_from_pdf: {str(e)}") @app.route('/') def index(): return render_template('app_index.html') # API endpoint @app.route('/process_pdf', methods=['POST']) def process_pdf(): try: logger.info("Received request to process PDF.") if 'pdf_file' not in request.files: logger.warning("No PDF file found in request.") return jsonify({"error": "Missing PDF file in form-data with key 'pdf_file'"}), 400 pdf_file = request.files['pdf_file'] if pdf_file.filename == '': return jsonify({"error": "Empty filename"}), 400 # Save the uploaded PDF temporarily filename = secure_filename(pdf_file.filename) temp_dir = tempfile.mkdtemp() saved_pdf_path = os.path.join(temp_dir, filename) pdf_file.save(saved_pdf_path) logger.info(f"Saved uploaded PDF to: {saved_pdf_path}") # Extract & process json_path = None output_path, result = extract_images_from_pdf(saved_pdf_path, json_path) logger.info("Received request to process PDF.") return jsonify({ "message": "✅ PDF processed successfully", "output_json": output_path, "sprites": result }) except Exception as e: logger.exception("❌ Failed to process PDF") return jsonify({"error": f"❌ Failed to process PDF: {str(e)}"}), 500 if __name__ == '__main__': app.run(host='0.0.0.0', port=7860, debug=True)