Spaces:
Build error
Build error
File size: 3,410 Bytes
533c554 d3cdd11 533c554 d3cdd11 533c554 d3cdd11 533c554 d3cdd11 533c554 d3cdd11 533c554 d3cdd11 533c554 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | from dotenv import load_dotenv
load_dotenv() # take environment variables from .env.
import streamlit as st
import os
from PIL import Image
import google.generativeai as genai
import json
os.getenv("GOOGLE_API_KEY")
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
## Function to load OpenAI model and get responses
def get_gemini_response(input, image, prompt):
model = genai.GenerativeModel('gemini-pro-vision')
response = model.generate_content([input, image[0], prompt])
return response.text
def input_image_setup(uploaded_file):
# Check if a file has been uploaded
if uploaded_file is not None:
# Read the file into bytes
bytes_data = uploaded_file.getvalue()
image_parts = [
{
"mime_type": uploaded_file.type, # Get the mime type of the uploaded file
"data": bytes_data
}
]
return image_parts
else:
return None # Return None if no file uploaded
# initialize our streamlit app
st.set_page_config(page_title="Crime Lens API")
st.header("Criminal Case Filing System OCR")
# Add option to use default FIR copy
use_default_copy = st.checkbox("Use Default FIR Copy")
uploaded_file = None
image = ""
if not use_default_copy:
uploaded_file = st.file_uploader("Choose an image...", type=["pdf", "jpg", "jpeg", "png"])
if use_default_copy or uploaded_file is not None:
if use_default_copy:
image_path = "fir-copy.jpg"
image = Image.open(image_path)
st.image(image, caption="Using Default FIR Copy", use_column_width=True)
else:
# User uploaded their own file
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image.", use_column_width=True)
submit = st.button("Extract data")
## If ask button is clicked
if submit:
input_prompt = """
Extract the data of FIR No, Date and Time of FIR, Occurrence of offence, Type of information, Place of occurrence, Complaint / Victim's name
"""
image_data = input_image_setup(uploaded_file) # Use uploaded file if available, otherwise None
if image_data is None: # Check if user selected default copy
# Use default image path if no file uploaded
image_data = [{'mime_type': 'image/jpeg', 'data': open('fir-copy.jpg', 'rb').read()}]
response = get_gemini_response(input_prompt, image_data, input_prompt)
st.subheader("The Response is")
print("Respons")
# Extracting field data
field_data = {
"FIR No": "",
"Date and Time of FIR": "",
"Occurrence of offence": "",
"Type of information": "",
"Place of occurrence": "",
"Victim": ""
}
# Parse response to extract field data
lines = response.split("\n")
for line in lines:
for field in field_data.keys():
if field in line:
value = line.split(":")[1].strip()
field_data[field] = value
# Display field data in a table
st.table(field_data)
# Save field data as JSON
json_filename = "field_data.json"
with open(json_filename, "w") as json_file:
json.dump(field_data, json_file)
# Add a download button for the JSON file
st.download_button(
label="Download JSON",
data=json.dumps(field_data),
file_name="field_data.json",
mime="application/json"
)
|