Hammad712 commited on
Commit
23c697b
·
verified ·
1 Parent(s): 96f5ea0

Rename app.py to main.py

Browse files
Files changed (2) hide show
  1. app.py +0 -59
  2. main.py +66 -0
app.py DELETED
@@ -1,59 +0,0 @@
1
- import streamlit as st
2
- import os
3
- import PIL.Image
4
- from pdf2image import convert_from_bytes
5
- from google import genai
6
- from google.genai import types
7
-
8
- # Retrieve the API key from the environment variable.
9
- API_KEY = os.getenv("API_KEY")
10
- if API_KEY is None:
11
- st.error("API Key not found. Please set the API_KEY environment variable.")
12
- st.stop()
13
-
14
- # Initialize the GenAI client.
15
- client = genai.Client(api_key=API_KEY)
16
-
17
- def extract_text_from_image(img):
18
- """
19
- Extracts text from a PIL image using the Google GenAI API.
20
- """
21
- response = client.models.generate_content(
22
- model="gemini-2.0-flash",
23
- contents=["Extract the text from the image. Do not write anything except the extracted content", img])
24
- return response.text
25
-
26
- st.title("PDF/Image Text Extraction using Google GenAI")
27
-
28
- # File uploader (accepts PDF and common image formats).
29
- uploaded_file = st.file_uploader("Upload a PDF or image file", type=["pdf", "png", "jpg", "jpeg", "webp"])
30
-
31
- if uploaded_file is not None:
32
- output_text = ""
33
- st.write("**Uploaded File:**", uploaded_file.name)
34
-
35
- # Process PDF files.
36
- if uploaded_file.name.lower().endswith(".pdf"):
37
- st.info("Processing PDF file...")
38
- with st.spinner("Converting PDF pages to images..."):
39
- images = convert_from_bytes(uploaded_file.read(), dpi=200)
40
-
41
- for idx, img in enumerate(images, start=1):
42
- with st.spinner(f"Extracting text from page {idx}..."):
43
- page_text = extract_text_from_image(img)
44
- output_text += f"### Page {idx}\n\n{page_text}\n\n"
45
- else:
46
- # Process image files.
47
- st.info("Processing image file...")
48
- with st.spinner("Extracting text from image..."):
49
- img = PIL.Image.open(uploaded_file)
50
- output_text += extract_text_from_image(img) + "\n\n"
51
-
52
- st.success("Extraction complete!")
53
-
54
- st.markdown("### Extracted Text:")
55
- st.markdown(output_text)
56
-
57
- # Provide a download button for the Markdown file.
58
- md_bytes = output_text.encode("utf-8")
59
- st.download_button(label="Download Markdown file", data=md_bytes, file_name="output.md", mime="text/markdown")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
main.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import io
3
+ import tempfile
4
+ import PIL.Image
5
+ from fastapi import FastAPI, File, UploadFile, HTTPException
6
+ from fastapi.responses import FileResponse
7
+ from pdf2image import convert_from_bytes
8
+ from google import genai
9
+ from google.genai import types
10
+
11
+ app = FastAPI(title="PDF/Image Text Extraction API")
12
+
13
+ # Retrieve the API key from an environment variable.
14
+ API_KEY = os.getenv("API_KEY")
15
+ if not API_KEY:
16
+ raise ValueError("API_KEY environment variable is not set")
17
+
18
+ # Initialize the GenAI client.
19
+ client = genai.Client(api_key=API_KEY)
20
+
21
+ def extract_text_from_image(img):
22
+ """
23
+ Extracts text from a PIL image using the Google GenAI API.
24
+ """
25
+ response = client.models.generate_content(
26
+ model="gemini-2.0-flash",
27
+ contents=["Extract the text from the image. Do not write anything except the extracted content", img]
28
+ )
29
+ return response.text
30
+
31
+ @app.post("/upload", summary="Upload a PDF or image file", response_description="Returns a Markdown file with the extracted text")
32
+ async def upload_file(file: UploadFile = File(...)):
33
+ if not file.filename:
34
+ raise HTTPException(status_code=400, detail="No file provided")
35
+
36
+ # Read file content.
37
+ file_contents = await file.read()
38
+ output_text = ""
39
+
40
+ if file.filename.lower().endswith(".pdf"):
41
+ try:
42
+ # Convert PDF bytes to images.
43
+ images = convert_from_bytes(file_contents, dpi=200)
44
+ except Exception as e:
45
+ raise HTTPException(status_code=500, detail=f"Error converting PDF: {str(e)}")
46
+
47
+ # Process each page.
48
+ for idx, img in enumerate(images, start=1):
49
+ page_text = extract_text_from_image(img)
50
+ output_text += f"### Page {idx}\n\n{page_text}\n\n"
51
+ else:
52
+ try:
53
+ # Process the file as an image.
54
+ img = PIL.Image.open(io.BytesIO(file_contents))
55
+ except Exception as e:
56
+ raise HTTPException(status_code=400, detail="Uploaded file is not a valid image")
57
+
58
+ output_text += extract_text_from_image(img) + "\n\n"
59
+
60
+ # Save the extracted text to a temporary Markdown file.
61
+ temp_md = tempfile.NamedTemporaryFile(delete=False, suffix=".md")
62
+ with open(temp_md.name, "w", encoding="utf-8") as md_file:
63
+ md_file.write(output_text)
64
+
65
+ # Return the file as a downloadable response.
66
+ return FileResponse(temp_md.name, filename="output.md", media_type="text/markdown")