Spaces:
Runtime error
Runtime error
File size: 5,328 Bytes
3f42a6f |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
import numpy as np
import cv2, csv, os
def mask_box(mask_img, points, color):
mask = np.ones_like(mask_img, dtype=np.uint8) * 255
cv2.fillPoly(mask, [points], color)
# Apply the mask to the image
img_with_overlay = cv2.bitwise_and(mask_img, mask)
return img_with_overlay
def mask_frame(mask_img, cl, tables, color):
# Blend the original image with the specified color
blend = 0.5
img_color = cv2.addWeighted(mask_img, blend, np.full_like(mask_img, color), 1 - blend, 0)
# Create a mask with a white rectangle in the specified region
mask = np.zeros_like(mask_img, dtype=np.uint8)
x1, y1, x2, y2 = cl.x, cl.y, cl.x + cl.w, cl.y + cl.h
cv2.rectangle(mask, (x1, y1), (x2, y2), (255, 255, 255), -1)
for table in tables:
for tab in table:
pts = np.array([(tab.x, tab.y), (tab.x+tab.w, tab.y), (tab.x+tab.w, tab.y+tab.h),(tab.x,tab.y+tab.h)], np.int32)
cv2.fillPoly(mask, [pts], (255, 255, 255))
# Apply the mask to choose between the blended image and the original
result = np.where(mask == 255, mask_img, img_color)
return result
def mask_img(img, gdt_boxes, tables, dimensions, frame, other_info):
mask_img=img.copy()
for table in tables:
for tab in table:
pts = np.array([(tab.x, tab.y), (tab.x+tab.w, tab.y), (tab.x+tab.w, tab.y+tab.h),(tab.x,tab.y+tab.h)], np.int32)
mask_img = mask_box(mask_img, pts, (212,242,247))
for gdt in gdt_boxes:
for g in gdt.values():
for tab in g:
pts = np.array([(tab.x, tab.y), (tab.x+tab.w, tab.y), (tab.x+tab.w, tab.y+tab.h),(tab.x,tab.y+tab.h)], np.int32)
mask_img = mask_box(mask_img, pts, (68,136,179))
if frame:
mask_img = mask_frame(mask_img, frame, tables, (100*2, 78*2, 73*2))
offset = (frame.x, frame.y)
else:
offset = (0, 0)
for dim in dimensions:
box = dim[1]
pts=np.array([(box[0]+offset),(box[1]+offset),(box[2]+offset),(box[3]+offset)])
mask_img = mask_box(mask_img, pts, (110,185,187))
for info in other_info:
box = info[1]
pts=np.array([(box[0]+offset),(box[1]+offset),(box[2]+offset),(box[3]+offset)])
mask_img = mask_box(mask_img, pts, (128,102,90))
return mask_img
def process_raw_output(output_path, table_results = None, gdt_results = None, dimension_results = None, other_info = None, save = False):
#Write Table Results
if table_results and save:
csv_file = os.path.join(output_path, 'table_results.csv')
new_table_results =[]
with open(csv_file, mode='w', newline='') as file:
writer = csv.writer(file)
# Write the header
for t in range(len(table_results)):
writer.writerow([f'Table_{t}'])
writer.writerow(["Text", "X Coordinate", "Y Coordinate"])
tab_results = []
for item in table_results[t]:
line = [item['text'], (item['left'], item['top'])]
tab_results.append(line)
writer.writerow([item['text'], item['left'], item['top']])
new_table_results.append(tab_results)
table_results = new_table_results
#Write GD&T Results
if gdt_results and save:
csv_file = os.path.join(output_path, 'gdt_results.csv')
with open(csv_file, mode='w', newline='') as file:
writer = csv.writer(file)
# Write the header
writer.writerow(["Text", "X Coordinate", "Y Coordinate"])
# Write the data
for item in gdt_results:
text, coords = item
writer.writerow([text, coords[0], coords[1]])
#Write Dimension Results
if dimension_results:
new_dim_results = []
for item in dimension_results:
text, coords = item
center = np.mean(coords, axis=0).astype(int).tolist()
new_dim_results.append([text, center])
dimension_results = new_dim_results
if save:
csv_file = os.path.join(output_path, 'dimension_results.csv')
with open(csv_file, mode='w', newline='') as file:
writer = csv.writer(file)
# Write the header
writer.writerow(["Text", "X Coordinate", "Y Coordinate"])
# Write the data
for i in new_dim_results:
writer.writerow([i[0], i[1][0], i[1][1]])
if other_info:
new_info_results = []
for item in other_info:
text, coords = item
center = np.mean(coords, axis=0).astype(int).tolist()
new_info_results.append([text, center])
other_info = new_info_results
if save:
csv_file = os.path.join(output_path, 'other_info.csv')
with open(csv_file, mode='w', newline='') as file:
writer = csv.writer(file)
# Write the header
writer.writerow(["Text", "X Coordinate", "Y Coordinate"])
# Write the data
for i in new_info_results:
writer.writerow([i[0], i[1][0], i[1][1]])
return table_results, gdt_results, dimension_results, other_info
|