app.py

import streamlit as st
import os
import torch
from paddleocr import PaddleOCR
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array, load_img
import numpy as np
from datetime import datetime, timedelta
import re
from ultralytics import YOLO
import pandas as pd

# Initialize models
ocr = PaddleOCR(lang='en')
fruit_model = load_model('DenseNet20_model.h5')
brand_model = YOLO('b.pt')

# Class names for fruit freshness classification
class_names = {
    0: 'Banana_Bad',
    1: 'Banana_Good',
    2: 'Fresh',
    3: 'FreshCarrot',
    4: 'FreshCucumber',
    5: 'FreshMango',
    6: 'FreshTomato',
    7: 'Guava_Bad',
    8: 'Guava_Good',
    9: 'Lime_Bad',
    10: 'Lime_Good',
    11: 'Rotten',
    12: 'RottenCarrot',
    13: 'RottenCucumber',
    14: 'RottenMango',
    15: 'RottenTomato',
    16: 'freshBread',
    17: 'rottenBread'
}

# Helper functions
def preprocess_image(image_path):
    img = load_img(image_path, target_size=(128, 128))
    img_array = img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = img_array / 255.0
    return img_array

def extract_expiry_dates(text):
    expiry_date_patterns = [
        r'USE BY (\d{1,2}[\/\-]\d{1,2}[\/\-]\d{4})',
        r'BEST BEFORE (\d{1,2}[\/\-]\d{1,2}[\/\-]\d{4})',
        r'(\d{1,2}[\/\-]\d{1,2}[\/\-]\d{4})',
        r'(\d{1,2}[\/\-]\d{1,2}[\/\-]\d{2})',
        r'(\d{1,2}\s*[A-Za-z]{3,}\s*\d{4})',
        r'(\d{1,2}\s*[A-Za-z]{3,}\s*\d{2})',
        r'(\d{4}[\/\-]\d{1,2}[\/\-]\d{1,2})',
        r'(\d{4}[A-Za-z]{3,}\d{1,2})',
        r'(\d{1,2}[A-Za-z]{3,}\d{4})',
        r'Best before (\d+) months'
    ]
    dates = []
    for pattern in expiry_date_patterns:
        matches = re.findall(pattern, text, re.IGNORECASE)
        dates.extend(matches)
    return dates

# Streamlit app
st.title("Image Processing Application")

# User choice for processing
task_choice = st.radio("Choose a task", ("Text and Brand Detection", "Freshness Detection"))

uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Ensure the uploads directory exists
    if not os.path.exists("uploads"):
        os.makedirs("uploads")

    # Save the uploaded file
    image_path = os.path.join("uploads", uploaded_file.name)
    with open(image_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    # Display the uploaded image
    st.image(image_path, caption='Uploaded Image', use_container_width=True)

    if task_choice == "Text and Brand Detection":
        # Text Extraction
        st.header("Text Extraction")
        result = ocr.ocr(image_path)
        text = ' '.join([line[1][0] for line in result[0]])
        st.write("Extracted Text:")
        st.text(text)

        # Expiry Date Extraction
        st.header("Expiry Date Extraction")
        expiry_dates = extract_expiry_dates(text)
        if expiry_dates:
            st.write("Expiry Dates Found:")
            for date in expiry_dates:
                st.text(date)
        else:
            st.write("No expiry dates found.")

        # Brand Prediction
        st.header("Brand Prediction")
        results = brand_model(image_path)
        detected_brands = []
        for result in results:
            boxes = result.boxes.data.cpu().numpy()
            for box in boxes:
                class_id = int(box[5])
                confidence = box[4]  # Assuming the confidence score is at index 4
                detected_brands.append((result.names[class_id], confidence))
        
        # Count occurrences and average confidence of each brand
        brand_data = {}
        for brand, confidence in detected_brands:
            if brand in brand_data:
                brand_data[brand]['count'] += 1
                brand_data[brand]['total_confidence'] += confidence
            else:
                brand_data[brand] = {'count': 1, 'total_confidence': confidence}
        
        # Prepare data for display
        brand_display_data = [
            {'Object': brand, 'Count': data['count'], 'Average Confidence': data['total_confidence'] / data['count']}
            for brand, data in brand_data.items()
        ]
        
        # Display detected brands in a table with column names
        st.write("Detected Brands and Counts:")
        st.table(pd.DataFrame(brand_display_data))

    elif task_choice == "Freshness Detection":
        # Freshness Prediction
        st.header("Freshness Prediction")
        img_array = preprocess_image(image_path)
        predictions = fruit_model.predict(img_array)
        predicted_class = np.argmax(predictions, axis=1)[0]
        st.write("Predicted Freshness:", class_names[predicted_class])