|
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') |
|
|
|
@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()) |
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
