Marthee commited on
Commit
e0ee72b
·
verified ·
1 Parent(s): dd0e903

Update doors_fasterrcnn.py

Browse files
Files changed (1) hide show
  1. doors_fasterrcnn.py +8 -7
doors_fasterrcnn.py CHANGED
@@ -109,7 +109,7 @@ def calculate_midpoint(top_left, bottom_right):
109
  ym = int((y1 + y2) / 2)
110
  return (xm, ym)
111
 
112
- def mid_points_bbox(bbox):
113
  midpoints = []
114
  for i in range(len(bbox)):
115
  x1 = int(bbox[i][0])
@@ -119,6 +119,7 @@ def mid_points_bbox(bbox):
119
  top_left_corner = (x1, y1)
120
  bottom_right_corner = (x2, y2)
121
  door_width = calculate_width(bbox[i])
 
122
  midpoint = calculate_midpoint(top_left_corner, bottom_right_corner)
123
  midpoints.append((midpoint, door_width))
124
  return midpoints
@@ -127,7 +128,7 @@ def create_annotations(door_kind, midpoints):
127
  door = door_kind
128
  annotations = []
129
  for i in range(len(midpoints)):
130
- annotations.append((midpoints[i][0][0],midpoints[i][0][1], door+f" with {midpoints[i][1]} pixels width"))
131
  return annotations
132
 
133
  def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d):
@@ -168,10 +169,10 @@ def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d):
168
 
169
  def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath): ####pdf_fullpath here is the data and not the path
170
  img_pillow = convert2pillow(pdf_fullpath)
171
- new_image = img_pillow.resize((2384, 1684))
172
 
173
  # Specify the number of classes (including the background)
174
- num_classes = 6 # Ensure this matches the saved model's number of classes
175
 
176
  # Load the model with the specified number of classes
177
  model = get_model(num_classes)
@@ -192,15 +193,15 @@ def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath): ####pdf_fullpath her
192
  model.to(device)
193
 
194
  # START INFERENCE
195
- sbox, dbox = ev_model(new_image, model, device, 0.6)
196
 
197
  #Dataframe for Doors count
198
  doors_count = {'Type': ['Single Doors', 'Double Doors'], 'Quantity': [len(sbox), len(dbox)]}
199
  df_doors = pd.DataFrame(doors_count)
200
 
201
  #Door's details for annotation
202
- single_info = mid_points_bbox(sbox)
203
- double_info = mid_points_bbox(dbox)
204
 
205
  #Annotations that will be printed on the pdf
206
  single_annotations = create_annotations("single door", single_info)
 
109
  ym = int((y1 + y2) / 2)
110
  return (xm, ym)
111
 
112
+ def mid_points_bbox(bbox, ratio):
113
  midpoints = []
114
  for i in range(len(bbox)):
115
  x1 = int(bbox[i][0])
 
119
  top_left_corner = (x1, y1)
120
  bottom_right_corner = (x2, y2)
121
  door_width = calculate_width(bbox[i])
122
+ door_width = door_width * ratio
123
  midpoint = calculate_midpoint(top_left_corner, bottom_right_corner)
124
  midpoints.append((midpoint, door_width))
125
  return midpoints
 
128
  door = door_kind
129
  annotations = []
130
  for i in range(len(midpoints)):
131
+ annotations.append((midpoints[i][0][0],midpoints[i][0][1], door+f" with {midpoints[i][1]} meters width"))
132
  return annotations
133
 
134
  def add_annotations_to_pdf(image, pdf_name, annotation_s, annotation_d):
 
169
 
170
  def main_run(pdf_fullpath, weights_path, pdf_name,pdfpath): ####pdf_fullpath here is the data and not the path
171
  img_pillow = convert2pillow(pdf_fullpath)
172
+ #new_image = img_pillow.resize((2384, 1684))
173
 
174
  # Specify the number of classes (including the background)
175
+ num_classes = 3 # Ensure this matches the saved model's number of classes
176
 
177
  # Load the model with the specified number of classes
178
  model = get_model(num_classes)
 
193
  model.to(device)
194
 
195
  # START INFERENCE
196
+ sbox, dbox = ev_model(img_pillow, model, device, 0.6)
197
 
198
  #Dataframe for Doors count
199
  doors_count = {'Type': ['Single Doors', 'Double Doors'], 'Quantity': [len(sbox), len(dbox)]}
200
  df_doors = pd.DataFrame(doors_count)
201
 
202
  #Door's details for annotation
203
+ single_info = mid_points_bbox(sbox, ratio)
204
+ double_info = mid_points_bbox(dbox, ratio)
205
 
206
  #Annotations that will be printed on the pdf
207
  single_annotations = create_annotations("single door", single_info)