Spaces:
No application file
No application file
| 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() | |