-
Notifications
You must be signed in to change notification settings - Fork 0
/
ocr.py
51 lines (47 loc) · 1.89 KB
/
ocr.py
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
import easyocr
import cv2
import os
from tqdm import tqdm
import numpy as np
import json
reader = easyocr.Reader(['en']) # this needs to run only once to load the model into memory
with open('data.json', 'r') as file:
data = json.load(file)
print(data[-1])
s = ""
t = data[-1]
#mixed_data/frames_Sequoia Capital
#count = 0
path = f"mixed_data/frames_{t['Author']}/{t['Title']}"
for img in tqdm(os.listdir(path), desc=f'Processing images for {t["Author"]}'):
img_path = os.path.join(path,img)
result = reader.readtext(img_path)
# Load the image
#image = cv2.imread("mixed_data/frames_Sequoia Capital/frame0049.png")
# Draw bounding boxes on the image
for box in result:
# Extract coordinates, label, and confidence
points, label, _ = box
s = s+label
# # Convert points to numpy array
# points = np.array(points)
# # Define the color of the rectangle in BGR format (Blue, Green, Red)
# color = (0, 255, 0) # Green color
# # Define the thickness of the rectangle border
# thickness = 2
# # Draw the rectangle on the image
# cv2.rectangle(image, (int(points[0][0]),int(points[0][1])), (int(points[2][0]),int(points[2][0])), color, thickness)
# # Draw the bounding box
# #cv2.polylines(image, [points], isClosed=True, color=(0, 255, 0), thickness=2)
# # Put label on the bounding box
# #cv2.putText(image, label, (points[0][0], points[0][1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# # # Display the image
# # cv2.imwrite(os.path.join("temp_output", f"image_with_bounding_boxes{count}.jpg"), image)
# # count+=1
print(s)
with open(f"mixed_data/output_{t['Author']}/" + f"ocr_{t['Title']}_text.txt", "w") as file:
file.write(s)
file.close()
# cv2.imshow("Image with Bounding Boxes", image)
# cv2.waitKey(0)
# cv2.destroyAllWindows()