app2.py

import os
import csv
import easyocr
import shutil
import random
import cv2
from glob import glob
from ultralytics import YOLOv10
import random
from glob import glob
from ultralytics import YOLOv10
import supervision as sva
from ultralytics import YOLOv10
import supervision as sv
import supervision as sv
from flask import Flask, request, jsonify, send_from_directory, render_template

import textwrap
app = Flask(__name__)

def enhance_contrast(image):
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    equalized_image = cv2.equalizeHist(gray_image)
    return equalized_image


def calculate_iou(bbox1, bbox2):
    x1_max = max(bbox1[0], bbox2[0])
    y1_max = max(bbox1[1], bbox2[1])
    x2_min = min(bbox1[2], bbox2[2])
    y2_min = min(bbox1[3], bbox2[3])

    inter_area = max(0, x2_min - x1_max) * max(0, y2_min - y1_max)

    bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
    bbox2_area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])

    iou = inter_area / float(bbox1_area + bbox2_area - inter_area) if (bbox1_area + bbox2_area - inter_area) > 0 else 0
    return iou


cropped_dir = "./app/cropped_images/"
if os.path.exists(cropped_dir):
        shutil.rmtree(cropped_dir)
os.makedirs(cropped_dir, exist_ok=True)

output_dir1 = "./app/Folder1"
output_dir2 = "./app/Folder2"
output_dir3 = "./app/Folder3"
UPLOAD_FOLDER = "./app/data1"
os.makedirs(output_dir1, exist_ok=True)
os.makedirs(output_dir2, exist_ok=True)
os.makedirs(output_dir3, exist_ok=True)
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

@app.route('/')
def index():
    return render_template('index3.html')  # This will serve your HTML page

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'invoice-upload' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['invoice-upload']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    if file:
        file_path = os.path.join(UPLOAD_FOLDER, file.filename)
        file.save(file_path)
        output_image, output_csv = process_image()

        return jsonify({
            'image_path': output_image,
            'csv_path': output_csv
        })

def process_image():
    print("Current working directory:", os.getcwd())
    
    # Check contents in the root directory
    print("Current directory contents:", os.listdir('/'))

    model = YOLOv10(f'./runs/detect/train3/weights/best (1).pt')
    dataset = sv.DetectionDataset.from_yolo(
        images_directory_path=f"./data/MyNewVersion5.0Dataset/valid/images",
        annotations_directory_path=f"./data/MyNewVersion5.0Dataset/valid/labels",
        data_yaml_path=f"./data/MyNewVersion5.0Dataset/data.yaml"
    )
    bounding_box_annotator = sv.BoundingBoxAnnotator()
    label_annotator = sv.LabelAnnotator()
    image_dir = "./app/data1"
    files = os.listdir('./app/data1')
    files.sort()
    files = files[0:100]
    print(files)
    counter = 0
    for ii in files:
        random_image_data = cv2.imread('./app/data1/' + ii)
        random_image_data1 = cv2.imread('./app/data1/' + ii)
        results = model(source='./app/data1/' + ii, conf=0.07)[0]
        detections = sv.Detections.from_ultralytics(results)
        annotated_image = bounding_box_annotator.annotate(scene=random_image_data, detections=detections)
        annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)
        save_path = "./app/Folder1/" + "detection" + ii
        cv2.imwrite(save_path, annotated_image)
        print(f"Annotated image saved at {save_path}")
        bounding_boxes = results.boxes.xyxy.cpu().numpy()
        class_ids = results.boxes.cls.cpu().numpy()
        confidences = results.boxes.conf.cpu().numpy()
        bounding_box_save_path = "./bounding_boxes.txt"
        with open(bounding_box_save_path, 'w') as f:
            for i, (bbox, class_id, confidence) in enumerate(zip(bounding_boxes, class_ids, confidences)):
                x1, y1, x2, y2 = map(int, bbox)
                f.write(f"Object {i + 1}: Class {class_id}, Confidence: {confidence:.2f}, "
                        f"Bounding box: ({x1}, {y1}, {x2}, {y2})\n")
                cropped_image = random_image_data1[y1:y2, x1:x2]
                cropped_image_path = os.path.join(cropped_dir, f"cropped_object_{i + 1}.jpg")
                cv2.imwrite(cropped_image_path, cropped_image)
                print(f"Enhanced cropped image saved at {cropped_image_path}")
        print(f"Checking contents of /app/data: {bounding_box_save_path}")
        print(f"Directory listing: {os.listdir('./app/Folder1')}")
        print(f"Bounding box coordinates saved at {bounding_box_save_path}")
        try:
          reader = easyocr.Reader(['en'],recog_network='en_sample',model_storage_directory='./EasyOCR-Trainer/EasyOCR/easyocr/model', user_network_directory='./EasyOCR-Trainer/EasyOCR/user_network')
        except Exception as e:
                print(f"Error initializing EasyOCR Reader: {e}")
                raise
        reader = easyocr.Reader(
            ['en'],
            recog_network='en_sample',
            model_storage_directory='./EasyOCR-Trainer/EasyOCR/easyocr/model',
            user_network_directory='./EasyOCR-Trainer/EasyOCR/user_network')
        import re
        input_file_path = './bounding_boxes.txt'
        cropped_images_folder = './app/cropped_images/'
        output_csv_path = './Folder2/' + ii + 'bounding_boxes_with_recognition.csv'
        print(f"Checking contents of ./app/data: {bounding_box_save_path}")
        print(f"Directory listing: {os.listdir('./app/data')}")

        with open(input_file_path, 'r') as infile:
            lines = infile.readlines()
        with open(output_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            csv_writer = csv.writer(csvfile)
            csv_writer.writerow(['Object ID', 'Bounding Box', 'Image Name', 'Recognized Text'])
            for i, line in enumerate(lines):
                object_id = f"Object_{i + 1}"
                bounding_box_info = line.strip()
                cropped_image_name = f"cropped_object_{i + 1}.jpg"
                cropped_image_path = os.path.join(cropped_images_folder, cropped_image_name)
                if os.path.exists(cropped_image_path):
                    bbox_match = re.search(r"Bounding box: \((\d+), (\d+), (\d+), (\d+)\)", bounding_box_info)
                    if bbox_match:
                        x1, y1, x2, y2 = map(int, bbox_match.groups())
                        detected_boxes = [[x1, x2, y1, y2]]
                    else:
                        print("No bounding box found in the info.")
                    cropped_image = cv2.imread(cropped_image_path, cv2.IMREAD_GRAYSCALE)
                    horizontal_list1, free_list1 = reader.detect(cropped_image)
                    free_list1 = free_list1 if free_list1 is not None else []
                    horizontal_list1 = [box for sublist in horizontal_list1 for box in sublist]
                    free_list1 = []
                    horizontal_list_for_recognize = detected_boxes if not horizontal_list1 else horizontal_list1
                    if horizontal_list1:
                        result = reader.recognize(cropped_image, detail=0, horizontal_list=horizontal_list1,
                                                  free_list=free_list1)
                    else:
                        result = reader.recognize(random_image_data1, detail=0, horizontal_list=detected_boxes,
                                                  free_list=free_list1)
                    recognized_text = ' '.join(result) if result else ''
                else:
                    recognized_text = 'No image found'
                csv_writer.writerow([object_id, bounding_box_info, cropped_image_name, recognized_text])
        print(f"CSV file with recognition results saved at {output_csv_path}")

        def calculate_iou(bbox1, bbox2):
            x1_max = max(bbox1[0], bbox2[0])
            y1_max = max(bbox1[1], bbox2[1])
            x2_min = min(bbox1[2], bbox2[2])
            y2_min = min(bbox1[3], bbox2[3])

            inter_area = max(0, x2_min - x1_max) * max(0, y2_min - y1_max)

            bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
            bbox2_area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])

            iou = inter_area / float(bbox1_area + bbox2_area - inter_area) if (bbox1_area + bbox2_area - inter_area) > 0 else 0
            return iou

        image_path = "/data1" + ii
        csv_file_path = output_csv_path = '/Folder2/' + ii + 'bounding_boxes_with_recognition.csv'
        image = cv2.imread(image_path)
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 1.3
        font_thickness = 2
        color = (255, 0, 255)
        bboxes = []
        recognized_texts = []
        with open(csv_file_path, 'r', encoding='utf-8') as csvfile:
            csv_reader = csv.DictReader(csvfile)
            for row in csv_reader:
                bbox_match = re.search(r'\((\d+), (\d+), (\d+), (\d+)\)', row['Bounding Box'])
                if bbox_match:
                    bbox = [int(bbox_match.group(i)) for i in range(1, 5)]
                    bboxes.append(bbox)
                    recognized_texts.append(row['Recognized Text'])
        filtered_bboxes = []
        filtered_texts = []
        iou_threshold = 0.4
        for i, bbox1 in enumerate(bboxes):
            keep = True
            for j, bbox2 in enumerate(filtered_bboxes):
                if calculate_iou(bbox1, bbox2) > iou_threshold:
                    keep = False
                    break
            if keep:
                filtered_bboxes.append(bbox1)
                filtered_texts.append(recognized_texts[i])
        for bbox, recognized_text in zip(filtered_bboxes, filtered_texts):
            x1, y1, x2, y2 = bbox
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            max_chars_per_line = 60
            wrapped_text = textwrap.wrap(recognized_text, width=max_chars_per_line)
            text_y = y1 - 10 if y1 - 10 > 10 else y1 + 10
            for line in wrapped_text:
                cv2.putText(image, line, (x1, text_y), font, font_scale, color, font_thickness)
                text_y += int(font_scale * 20)
            output_image_path = "/Folder3/" + "annotated" + ii + ".png"
            cv2.imwrite(output_image_path, image)
            print(f"Annotated image saved at {output_image_path}")
            counter += 1

@app.route('/download_csv/<filename>')
def download_csv(filename):
    return send_from_directory(output_dir2, filename, as_attachment=True)

@app.route('/download_image/<filename>')
def download_image(filename):
    return send_from_directory(output_dir3, filename, as_attachment=True)