Spaces:
No application file
No application file
File size: 2,870 Bytes
4c1b733 |
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 |
import streamlit as st
from ultralytics import YOLO
from PIL import Image, ImageDraw
import tempfile
def greeting(name):
return f"Hello from module2, {name}!"
# Function to process the image, classify it, and crop if clear
def process_image(file_path):
model = YOLO("best.pt", "v8")
# Predict with the model
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_image:
temp_image.close()
Image.open(file_path).save(temp_image.name)
results = model.predict(source=temp_image.name, conf=0.4, save=True)
blur_conf_threshold = 0.5
clear_conf_threshold = 0.9
# Initialize flags
is_blur = False
is_clear = False
cropped_image = None
# Process results
for result in results[0].boxes:
confidence = result.conf[0].item() # Extract the confidence score
if blur_conf_threshold <= confidence <= clear_conf_threshold:
is_blur = True
elif confidence > clear_conf_threshold:
is_clear = True
box = result.xyxy[0].tolist() # Extract bounding box coordinates
cropped_image = crop_image(file_path, box)
# Return classification and cropped image
if is_blur:
return 'The image is blurry. Please reupload the image again!', None
elif is_clear:
return 'The image is clear', cropped_image
else:
s=("Not Detected! The image is uncertain. Please reupload the image again!")
return s, None
# Function to crop the image based on bounding box
def crop_image(file_path, box):
image = Image.open(file_path)
cropped_image = image.crop(box)
return cropped_image
def greeting(name):
return f"Hello module2, {name}!"
# Streamlit app
def main():
st.title('Welcome to my AI project')
st.title('Document Detection')
st.text('This is a web app to:\n1- Detect documents\n2- Classify if document is clear or blurry\n3- Crop the document image!')
# File uploader
uploaded_file = st.file_uploader('Upload your image here:', type=['png', 'jpg', 'jpeg'])
if uploaded_file is not None:
# Display the uploaded image
image = Image.open(uploaded_file)
st.image(image, caption='Uploaded Image.', use_column_width=True)
st.success("Photo uploaded successfully!")
# Process the image and get classification and cropped image
classification, cropped_image = process_image(uploaded_file)
# Display classification result
st.write('Classification result:', classification)
# Display cropped image if classification is clear
if cropped_image is not None:
st.image(cropped_image, caption='Cropped Document Image.', use_column_width=True)
if __name__ == "__main__":
main()
|