car-view-classifier / app /dependencies /yolo_classification_6.py
jayvatliq's picture
Upload folder using huggingface_hub
b1adb1e verified
Raw
History Blame Contribute Delete
32.7 kB
import os
import shutil
from datetime import datetime
import cv2
import matplotlib.pyplot as plt
import numpy as np
from fastapi import UploadFile
from PIL import Image
from ultralytics import YOLO
# Models loaded once at import time
car_detection_model = YOLO(r"car_detection.pt")
part_detection_model = YOLO(r"car_part.pt")
DRIVER_SIDE = "Driver Side View"
PASSENGER_SIDE = "Passenger Side View"
FRONT_DRIVER_CORNER = "Front Driver Side Corner View"
FRONT_PASSENGER_CORNER = "Front Passenger Side Corner View"
REAR_DRIVER_CORNER = "Rear Driver Side Corner View"
REAR_PASSENGER_CORNER = "Rear Passenger Side Corner View"
CORNER_SWAP = {
FRONT_PASSENGER_CORNER: FRONT_DRIVER_CORNER,
FRONT_DRIVER_CORNER: FRONT_PASSENGER_CORNER,
REAR_PASSENGER_CORNER: REAR_DRIVER_CORNER,
REAR_DRIVER_CORNER: REAR_PASSENGER_CORNER,
}
def _enforce_side(label, desired_side):
"""Return label with the correct driver/passenger side applied."""
if not desired_side or not label or label == "NA":
return label
if label in (DRIVER_SIDE, PASSENGER_SIDE):
return desired_side if label != desired_side else label
if label in CORNER_SWAP:
wrong_keyword = "Passenger" if desired_side == DRIVER_SIDE else "Driver"
return CORNER_SWAP[label] if wrong_keyword in label else label
return label
MIN_RATIO = {
# Front
"Front-bumper": 0.015,
"Grille": 0.010,
"Headlight": 0.005, # small but important
"Hood": 0.020,
"License-plate": 0.002, # very small
# Side
"Front-door": 0.030,
"Back-door": 0.030,
"Front-wheel": 0.010,
"Back-wheel": 0.010,
"Mirror": 0.001, # small but always relevant
"Quarter-panel": 0.010,
"Rocker-panel": 0.010,
"Roof": 0.020,
# Windows
"Windshield": 0.020,
"Front-window": 0.015,
"Back-window": 0.015,
"Back-windshield": 0.020,
# Rear
"Back-bumper": 0.015,
"Tail-light": 0.005, # small but important
"Trunk": 0.020,
# Catch-all fallback
"default": 0.01
}
viewing_angle_rules = {
"Front Right": {
"must_be_visible": ["Front-bumper", "Grille", "Headlight", "Front-wheel", "Windshield", "Front-door", "Front-window", "Fender", "Mirror", "Rocker-panel"],
"optional_parts": ["Hood", "Roof", "Back-door", "Back-wheel", "Back-window", "Quarter-panel"],
"conflict_parts": ["Tail-light", "Back-bumper", "Back-windshield", "Trunk"]
},
"Right": {
"must_be_visible": ["Front-door", "Back-door", "Mirror", "Quarter-panel", "Fender", "Rocker-panel", "Front-wheel", "Back-wheel", "Back-window", "Front-window"],
"optional_parts": ["Roof", "Front-bumper", "Back-bumper", "Headlight", "Tail-light", "Hood", "Back-windshield", "Windshield", "Trunk", "Grille", "License-plate"],
"conflict_parts": []
},
"Rear Right": {
"must_be_visible": ["Back-bumper", "Tail-light", "Back-wheel", "Back-door", "Back-window", "Quarter-panel", "Back-windshield", "Rocker-panel", "Trunk"],
"optional_parts": ["Roof", "License-plate", "Front-wheel", "Front-door", "Fender", "Mirror", "Front-window"],
"conflict_parts": ["Front-bumper", "Headlight", "Grille", "Windshield", "Hood"]
},
"Rear Left": {
"must_be_visible": ["Back-bumper", "Tail-light", "Back-wheel", "Back-door", "Back-window", "Quarter-panel", "Back-windshield", "Rocker-panel", "Trunk"],
"optional_parts": ["Roof", "License-plate", "Front-wheel", "Front-door", "Fender", "Mirror", "Front-window"],
"conflict_parts": ["Front-bumper", "Headlight", "Grille", "Windshield", "Hood"]
},
"Left": {
"must_be_visible": ["Front-door", "Back-door", "Mirror", "Quarter-panel", "Fender", "Rocker-panel", "Front-wheel", "Back-wheel", "Back-window", "Front-window"],
"optional_parts": ["Roof", "Front-bumper", "Back-bumper", "Headlight", "Tail-light", "Hood", "Back-windshield", "Windshield", "Trunk", "Grille", "License-plate"],
"conflict_parts": []
},
"Front Left": {
"must_be_visible": ["Front-bumper", "Grille", "Headlight", "Front-wheel", "Windshield", "Front-door", "Front-window", "Fender", "Mirror", "Rocker-panel"],
"optional_parts": ["Hood", "Roof", "Back-door", "Back-wheel", "Back-window", "Quarter-panel"],
"conflict_parts": ["Tail-light", "Back-bumper", "Back-windshield", "Trunk"]
}
}
def compute_direction_mirror_refined(detected, fixed_label="Mirror"):
if fixed_label not in detected:
return "Unknown", None, None
mirror_box = detected[fixed_label][0]
mirror_center = ((mirror_box[0] + mirror_box[2]) / 2, (mirror_box[1] + mirror_box[3]) / 2)
groupA = ["Windshield", "Hood", "Headlight", "Front-bumper", "Front-wheel"]
groupB = ["Back-wheel", "Back-door", "Quarter-panel", "Rocker-panel", "Back-window"]
offsets_A = []
offsets_B = []
radial_lines = []
for part in groupA:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - mirror_center[0]
offsets_A.append(offset)
radial_lines.append((mirror_center, part_center))
for part in groupB:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - mirror_center[0]
offsets_B.append(offset)
radial_lines.append((mirror_center, part_center))
voteA = None
voteB = None
if offsets_A:
avg_A = sum(offsets_A) / len(offsets_A)
voteA = "Right" if avg_A > 0 else "Left"
if offsets_B:
avg_B = sum(offsets_B) / len(offsets_B)
voteB = "Right" if avg_B < 0 else "Left"
if voteA and voteB:
if voteA == voteB:
direction = f"{voteA} side view"
else:
if len(offsets_A) > len(offsets_B):
direction = f"{voteA} side view"
elif len(offsets_B) > len(offsets_A):
direction = f"{voteB} side view"
else:
direction = f"{voteB} side view"
elif voteA:
direction = f"{voteA} side view"
elif voteB:
direction = f"{voteB} side view"
else:
direction = "Unknown"
if voteA and voteB:
consensus = (voteA == voteB)
elif voteA or voteB:
consensus = True
else:
consensus = None
return direction, radial_lines, consensus
def compute_direction_front_wheel_refined(detected, fixed_label="Front-wheel"):
if fixed_label not in detected:
return "Unknown", None, None
front_wheel_box = detected[fixed_label][0]
front_wheel_center = ((front_wheel_box[0] + front_wheel_box[2]) / 2, (front_wheel_box[1] + front_wheel_box[3]) / 2)
groupA = ["Front-bumper", "Headlight", "Fender", "Grille"]
groupB = ["Front-door", "Back-door", "Rocker-panel", "Front-window", "Back-window", "Back-wheel", "Quarter-panel", "Mirror", "Windshield"]
offsets_A = []
offsets_B = []
radial_lines = []
for part in groupA:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - front_wheel_center[0]
offsets_A.append(offset)
radial_lines.append((front_wheel_center, part_center))
for part in groupB:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - front_wheel_center[0]
offsets_B.append(offset)
radial_lines.append((front_wheel_center, part_center))
voteA = None
voteB = None
if offsets_A:
avg_A = sum(offsets_A) / len(offsets_A)
voteA = "Right" if avg_A > 0 else "Left"
if offsets_B:
avg_B = sum(offsets_B) / len(offsets_B)
voteB = "Right" if avg_B < 0 else "Left"
if voteA and voteB:
if voteA == voteB:
direction = f"{voteA} side view"
else:
if len(offsets_A) > len(offsets_B):
direction = f"{voteA} side view"
elif len(offsets_B) > len(offsets_A):
direction = f"{voteB} side view"
else:
direction = f"{voteB} side view"
elif voteA:
direction = f"{voteA} side view"
elif voteB:
direction = f"{voteB} side view"
else:
direction = "Unknown"
if voteA and voteB:
consensus = (voteA == voteB)
elif voteA or voteB:
consensus = True
else:
consensus = None
return direction, radial_lines, consensus
def compute_direction_back_wheel_refined(detected, fixed_label="Back-wheel"):
if fixed_label not in detected:
return "Unknown", None, None
back_wheel_box = detected[fixed_label][0]
back_wheel_center = ((back_wheel_box[0] + back_wheel_box[2]) / 2,
(back_wheel_box[1] + back_wheel_box[3]) / 2)
groupA = ["Back-bumper", "Tail-light", "Quarter-panel", "Trunk"]
groupB = ["Front-wheel", "Rocker-panel", "Back-door", "Front-door",
"Back-window", "Front-window", "Mirror", "Fender"]
offsets_A = []
offsets_B = []
radial_lines = []
for part in groupA:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - back_wheel_center[0]
offsets_A.append(offset)
radial_lines.append((back_wheel_center, part_center))
for part in groupB:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - back_wheel_center[0]
offsets_B.append(offset)
radial_lines.append((back_wheel_center, part_center))
voteA = None
voteB = None
if offsets_A:
avg_A = sum(offsets_A) / len(offsets_A)
voteA = "Right" if avg_A < 0 else "Left"
if offsets_B:
avg_B = sum(offsets_B) / len(offsets_B)
voteB = "Left" if avg_B < 0 else "Right"
if voteA and voteB:
if voteA == voteB:
direction = f"{voteA} side view"
else:
if len(offsets_A) > len(offsets_B):
direction = f"{voteA} side view"
elif len(offsets_B) > len(offsets_A):
direction = f"{voteB} side view"
else:
direction = f"{voteB} side view"
elif voteA:
direction = f"{voteA} side view"
elif voteB:
direction = f"{voteB} side view"
else:
direction = "Unknown"
if voteA and voteB:
consensus = (voteA == voteB)
elif voteA or voteB:
consensus = True
else:
consensus = None
return direction, radial_lines, consensus
def compute_direction_headlight_refined(detected, fixed_label="Headlight"):
if fixed_label not in detected:
return "Unknown", None, None
headlight_box = detected[fixed_label][0]
headlight_center = ((headlight_box[0] + headlight_box[2]) / 2, (headlight_box[1] + headlight_box[3]) / 2)
groupA = ["Front-wheel", "Fender", "Mirror", "Rocker-panel", "Front-door", "Front-window"]
groupB = ["Front-bumper", "Grille", "Hood", "Windshield"]
offsets_A = []
offsets_B = []
radial_lines = []
for part in groupA:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - headlight_center[0]
offsets_A.append(offset)
radial_lines.append((headlight_center, part_center))
for part in groupB:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - headlight_center[0]
offsets_B.append(offset)
radial_lines.append((headlight_center, part_center))
voteA = None
voteB = None
if offsets_A:
avg_A = sum(offsets_A) / len(offsets_A)
voteA = "Right" if avg_A < 0 else "Left"
if offsets_B:
avg_B = sum(offsets_B) / len(offsets_B)
voteB = "Right" if avg_B > 0 else "Left"
if voteA and voteB:
if voteA == voteB:
direction = f"{voteA} side view"
else:
if len(offsets_A) > len(offsets_B):
direction = f"{voteA} side view"
elif len(offsets_B) > len(offsets_A):
direction = f"{voteB} side view"
else:
direction = f"{voteB} side view"
elif voteA:
direction = f"{voteA} side view"
elif voteB:
direction = f"{voteB} side view"
else:
direction = "Unknown"
if voteA and voteB:
consensus = (voteA == voteB)
elif voteA or voteB:
consensus = True
else:
consensus = None
return direction, radial_lines, consensus
def compute_direction_tail_refined(detected, fixed_label="Tail-light"):
if fixed_label not in detected:
return "Unknown", None, None
tail_box = detected[fixed_label][0]
tail_center = ((tail_box[0] + tail_box[2]) / 2, (tail_box[1] + tail_box[3]) / 2)
groupA = ["Trunk", "Back-bumper", "Back-windshield"]
groupB = ["Back-wheel", "Quarter-panel", "Back-door", "Back-window"]
offsets_A = []
offsets_B = []
radial_lines = []
for part in groupA:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - tail_center[0]
offsets_A.append(offset)
radial_lines.append((tail_center, part_center))
for part in groupB:
if part in detected:
for box in detected[part]:
part_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
offset = part_center[0] - tail_center[0]
offsets_B.append(offset)
radial_lines.append((tail_center, part_center))
voteA = None
voteB = None
if offsets_A:
avg_A = sum(offsets_A) / len(offsets_A)
voteA = "Right" if avg_A < 0 else "Left"
if offsets_B:
avg_B = sum(offsets_B) / len(offsets_B)
voteB = "Right" if avg_B > 0 else "Left"
if voteA and voteB:
if voteA == voteB:
direction = f"{voteA} side view"
else:
if len(offsets_A) > len(offsets_B):
direction = f"{voteA} side view"
elif len(offsets_B) > len(offsets_A):
direction = f"{voteB} side view"
else:
direction = f"{voteB} side view"
elif voteA:
direction = f"{voteA} side view"
elif voteB:
direction = f"{voteB} side view"
else:
direction = "Unknown"
if voteA and voteB:
consensus = (voteA == voteB)
elif voteA or voteB:
consensus = True
else:
consensus = None
return direction, radial_lines, consensus
def determine_vehicle_directions(detected):
directions = {}
radial_lines_all = {}
consensus_flags = {}
mirror_direction, mirror_radials, mirror_consensus = compute_direction_mirror_refined(detected, fixed_label="Mirror")
directions["Mirror"] = mirror_direction
radial_lines_all["Mirror"] = mirror_radials
consensus_flags["Mirror"] = mirror_consensus
tail_direction, tail_radials, tail_consensus = compute_direction_tail_refined(detected, fixed_label="Tail-light")
directions["Tail-light"] = tail_direction
radial_lines_all["Tail-light"] = tail_radials
consensus_flags["Tail-light"] = tail_consensus
front_wheel_direction, front_wheel_radials, front_wheel_consensus = compute_direction_front_wheel_refined(detected, fixed_label="Front-wheel")
directions["Front-wheel"] = front_wheel_direction
radial_lines_all["Front-wheel"] = front_wheel_radials
consensus_flags["Front-wheel"] = front_wheel_consensus
back_wheel_direction, back_wheel_radials, back_wheel_consensus = compute_direction_back_wheel_refined(detected, fixed_label="Back-wheel")
directions["Back-wheel"] = back_wheel_direction
radial_lines_all["Back-wheel"] = back_wheel_radials
consensus_flags["Back-wheel"] = back_wheel_consensus
headlight_direction, headlight_radials, headlight_consensus = compute_direction_headlight_refined(detected, fixed_label="Headlight")
directions["Headlight"] = headlight_direction
radial_lines_all["Headlight"] = headlight_radials
consensus_flags["Headlight"] = headlight_consensus
return directions, radial_lines_all, consensus_flags
def find_best_combination(pil_image):
"""
Runs car + part detection on pil_image.
Returns (pil_image, detected_parts_dict, detections_dict).
Returns (pil_image, {}, {}) if no car is found.
"""
image_array = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
car_results = car_detection_model(image_array)
car_detections_to_plot = []
main_car_bbox = None
best_area = 0
for result in car_results:
for box in result.boxes.data:
conf = float(box[4])
if conf < 0.50:
continue
x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
car_detections_to_plot.append((x1, y1, x2, y2, "Car", conf))
area = (x2 - x1) * (y2 - y1)
if area > best_area:
best_area = area
main_car_bbox = (x1, y1, x2, y2)
if main_car_bbox is None:
return pil_image, {}, {}
results = part_detection_model(image_array)
part_detections_to_plot = []
detected = {}
car_area = (main_car_bbox[2] - main_car_bbox[0]) * (main_car_bbox[3] - main_car_bbox[1])
for result in results:
for box in result.boxes.data:
conf = float(box[4])
if conf < 0.65:
continue
x1, y1, x2, y2, conf, class_id = box
label = result.names[int(class_id)]
width = max(0, x2 - x1)
height = max(0, y2 - y1)
area = width * height
part_ratio = area / car_area
min_ratio = MIN_RATIO.get(label, MIN_RATIO["default"])
if part_ratio < min_ratio and conf < 0.8:
continue
margin = 5
if (x1 <= main_car_bbox[0] + margin or x2 >= main_car_bbox[2] - margin or
y1 <= main_car_bbox[1] + margin or y2 >= main_car_bbox[3] - margin):
if part_ratio < min_ratio * 3:
continue
detected.setdefault(label, []).append((x1, y1, x2, y2))
part_detections_to_plot.append((x1, y1, x2, y2, label, conf))
detections = {
"car_detection": {
"image_array": image_array,
"detections_to_plot": car_detections_to_plot,
"title": "Car Detections",
},
"part_detection": {
"image_array": image_array,
"detections_to_plot": part_detections_to_plot,
"title": "Part Detections",
},
}
return pil_image, detected, detections
def determine_viewing_angle(detected):
need_review = False
detected_parts = set(detected.keys())
angle_scores = []
critical_parts = {
"Front Right": {"Front-bumper", "Headlight", "Front-door"},
"Right": {"Front-door", "Back-door", "Front-wheel", "Back-wheel"},
"Rear Right": {"Tail-light", "Trunk", "Back-windshield", "Back-door"},
"Rear Left": {"Tail-light", "Trunk", "Back-windshield", "Back-door"},
"Left": {"Front-door", "Back-door", "Front-wheel", "Back-wheel"},
"Front Left": {"Front-bumper", "Headlight", "Front-door"}
}
for angle, rules in viewing_angle_rules.items():
if angle in critical_parts:
total_critical = len(critical_parts[angle])
detected_critical = len(critical_parts[angle].intersection(detected_parts))
critical_ratio = detected_critical / total_critical
else:
critical_ratio = None
ess_score = sum(3 for part in rules['must_be_visible'] if part in detected_parts)
opt_score = sum(1 for part in rules['optional_parts'] if part in detected_parts)
conf_pen = sum(-3 for part in rules['conflict_parts'] if part in detected_parts)
raw_score = ess_score + opt_score + conf_pen
total_defined = len(rules['must_be_visible']) + len(rules['optional_parts']) + len(rules['conflict_parts'])
stage2_score = raw_score / total_defined if total_defined > 0 else 0.0
angle_scores.append((angle, critical_ratio, stage2_score))
print(angle_scores)
score_map = {angle.lower(): score for (angle, _, score) in angle_scores}
eps = 1e-6
stage_2_thresold = 0.75
sorted_all = sorted(angle_scores, key=lambda x: x[2], reverse=True)
symmetric_view = {
"Front Right": "Front Left",
"Front Left": "Front Right",
"Right": "Left",
"Left": "Right",
"Rear Right": "Rear Left",
"Rear Left": "Rear Right"
}
top_2_predictions = ["", ""]
if sorted_all:
top1 = sorted_all[0]
top1_key = top1[0].lower()
top1_score = top1[2]
top_2_predictions[0] = top1_key
second_key = ""
for angle, _, score in sorted_all[1:]:
if score <= top1_score and score >= stage_2_thresold and symmetric_view.get(angle, "").lower() != top1_key:
second_key = angle.lower()
break
if top1_score > 0.25 and score / (top1_score - eps) >= 0.60 and symmetric_view.get(angle, "").lower() != top1_key:
second_key = angle.lower()
break
top_2_predictions[1] = second_key
angles_above_60 = [item for item in angle_scores if item[1] is not None and item[1] >= 0.60]
if len(angles_above_60) >= 3:
need_review = True
best_angle, *_ = max(angle_scores, key=lambda x: x[2])
else:
candidates = [item for item in angle_scores if item[1] is not None and item[1] >= 0.9]
if candidates:
best_angle, *_ = max(candidates, key=lambda x: x[1])
else:
best_angle, *_ = max(angle_scores, key=lambda x: x[2])
if best_angle in ["Front", "Rear"]:
directions = {"Selected": best_angle}
else:
directions_all, _, consensus_all = determine_vehicle_directions(detected)
print(directions_all)
consensus_votes = {
ref: directions_all[ref]
for ref, flag in consensus_all.items()
if flag is True and directions_all[ref] != "Unknown"
}
if consensus_votes:
_, chosen_dir = next(iter(consensus_votes.items()))
majority_side = chosen_dir.split()[0]
else:
votes = {}
weights = {"Front-wheel": 1, "Back-wheel": 1, "Mirror": 1, "Headlight": 0.5, "Tail-light": 0.5}
for ref, dir_val in directions_all.items():
if dir_val != "Unknown":
side = dir_val.split()[0]
weight = weights.get(ref, 1)
votes[side] = votes.get(side, 0) + weight
majority_side = max(votes, key=votes.get) if votes else "Unknown"
if best_angle.startswith("Front"):
final_side_classification = "Front " + majority_side
elif best_angle.startswith("Rear"):
final_side_classification = "Rear " + majority_side
else:
final_side_classification = majority_side
directions = {"Selected": final_side_classification}
return directions, detected_parts, need_review, top_2_predictions, score_map
def deskew_image(pil_image: Image.Image) -> Image.Image:
img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
lines = cv2.HoughLines(edges, 1, np.pi / 180, 200)
angles = []
if lines is not None:
for _, theta in lines[:, 0]:
angle = (theta * 180 / np.pi) - 90
if angle < -90:
angle += 180
if angle > 90:
angle -= 180
angles.append(angle)
median_angle = np.median(angles) if len(angles) > 0 else 0
(h, w) = img.shape[:2]
center = (w // 2, h // 2)
M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
rotated = cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
return Image.fromarray(cv2.cvtColor(rotated, cv2.COLOR_BGR2RGB))
async def yolo_rule_based_classification(pil_image, image_name, img_file: UploadFile):
"""
Returns (final_secondaries, final_review, final_scores, file_details).
Tries each rotation (0, 90, -90) and also a deskewed variant — picks the
image variant that yields the most detected parts.
"""
rotations = [0, 90, -90]
best_raw_direction = None
final_review = False
fail_count = 0
best_top_predictions = [False, False]
best_score_map = {}
max_detected_labels = 0
detected_labels = {}
final_detections = {}
for rotation in rotations:
rotated_image = pil_image.rotate(rotation, expand=True) if rotation != 0 else pil_image
# Try original orientation
_, detected, detections = find_best_combination(rotated_image)
print(detected)
if len(detected.keys()) > max_detected_labels:
max_detected_labels = len(detected.keys())
detected_labels = detected
final_detections = detections
# Try deskewed version of this rotation
deskewed_image = deskew_image(rotated_image)
_, detected, detections = find_best_combination(deskewed_image)
print(detected)
if len(detected.keys()) > max_detected_labels:
max_detected_labels = len(detected.keys())
detected_labels = detected
final_detections = detections
print(detected_labels)
direction, detected_labels, need_review, top_predictions, score_map = determine_viewing_angle(detected_labels)
best_raw_direction = direction['Selected']
final_review = need_review
best_top_predictions = top_predictions
best_score_map = score_map
viewing_angle_map = {
"front": "Front View",
"rear": "Rear View",
"left": DRIVER_SIDE,
"right": PASSENGER_SIDE,
"left side view": DRIVER_SIDE,
"right side view": PASSENGER_SIDE,
"front right": FRONT_PASSENGER_CORNER,
"front left": FRONT_DRIVER_CORNER,
"rear right": REAR_PASSENGER_CORNER,
"rear left": REAR_DRIVER_CORNER,
"unknown": "NA",
}
desired_side = None
if isinstance(best_raw_direction, str):
raw = best_raw_direction.lower()
if "left" in raw:
desired_side = "Driver Side View"
elif "right" in raw:
desired_side = "Passenger Side View"
stage1_top1_key = best_top_predictions[0] if isinstance(best_top_predictions, (list, tuple)) and len(best_top_predictions) > 0 else False
stage1_top2_key = best_top_predictions[1] if isinstance(best_top_predictions, (list, tuple)) and len(best_top_predictions) > 1 else False
mapped_primary = "NA"
if stage1_top1_key and isinstance(stage1_top1_key, str):
mapped_primary = viewing_angle_map.get(stage1_top1_key.lower(), "NA")
elif isinstance(best_raw_direction, str):
mapped_primary = viewing_angle_map.get(best_raw_direction.lower(), "NA")
mapped_primary = _enforce_side(mapped_primary, desired_side)
final_secondaries = [False, False]
threshold = 0.75
eps = 1e-6
if stage1_top1_key and isinstance(stage1_top1_key, str):
mapped_top1 = viewing_angle_map.get(stage1_top1_key.lower(), "NA")
else:
mapped_top1 = mapped_primary if mapped_primary != "NA" else False
mapped_top1 = _enforce_side(mapped_top1, desired_side)
final_secondaries[0] = mapped_top1 if mapped_top1 != "NA" else False
mapped_top2 = False
if isinstance(stage1_top2_key, str):
top2_score = best_score_map.get(stage1_top2_key.lower(), -999.0)
top1_score = best_score_map.get(stage1_top1_key.lower(), -999.0) if isinstance(stage1_top1_key, str) else -999.0
if (top2_score >= threshold and top2_score <= top1_score) or (top2_score / (top1_score - eps) >= 0.60 and top1_score > 0.25):
mapped_top2 = _enforce_side(viewing_angle_map.get(stage1_top2_key.lower(), "NA"), desired_side)
else:
mapped_top2 = False
final_secondaries[1] = mapped_top2 if mapped_top2 and mapped_top2 != "NA" else False
if mapped_primary == "NA" and final_secondaries[0]:
mapped_primary = final_secondaries[0]
if fail_count >= 3:
final_review = True
final_scores = [0, 0]
if isinstance(stage1_top1_key, str):
final_scores[0] = best_score_map.get(stage1_top1_key.lower(), 0)
if isinstance(stage1_top2_key, str) and final_secondaries[1]:
final_scores[1] = best_score_map.get(stage1_top2_key.lower(), 0)
final_secondaries = [item for item in final_secondaries if isinstance(item, str)]
print(final_secondaries)
file_details = await store_images(final_detections, image_name, img_file)
return final_secondaries, final_review, final_scores, file_details
async def store_images(final_detections, image_name, img_file):
os.makedirs("./output_files", exist_ok=True)
img_folder, extension = os.path.splitext(image_name)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
ff = f"{img_folder}_{timestamp}"
output_folder = f"./output_files/{ff}"
os.makedirs(output_folder, exist_ok=True)
output_file_car = f"car_detection{extension}"
output_file_part = f"part_detection{extension}"
main_img_path = os.path.join(output_folder, img_file.filename)
await img_file.seek(0)
with open(main_img_path, "wb") as buffer:
shutil.copyfileobj(img_file.file, buffer)
try:
save_detection_img(
final_detections["car_detection"]["image_array"],
final_detections["car_detection"]["detections_to_plot"],
output_folder,
output_file_car,
final_detections["car_detection"]["title"],
)
except Exception as e:
print("Error Saving Car detections", e)
try:
save_detection_img(
final_detections["part_detection"]["image_array"],
final_detections["part_detection"]["detections_to_plot"],
output_folder,
output_file_part,
final_detections["part_detection"]["title"],
)
except Exception as e:
print("Error Saving Part detections", e)
return {
"main_img_name": f"{ff}/{img_file.filename}",
"part_detection": f"{ff}/{output_file_part}",
"car_detection": f"{ff}/{output_file_car}",
}
def save_detection_img(image_array, detections, save_folder, filename, title="Detections"):
vis_img = image_array.copy()
for x1, y1, x2, y2, label, conf in detections:
cv2.rectangle(vis_img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 3)
text = f"{label} {conf:.2f}"
text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
cv2.rectangle(
vis_img,
(int(x1), int(y1) - text_size[1] - 10),
(int(x1) + text_size[0], int(y1)),
(0, 255, 0),
-1,
)
cv2.putText(vis_img, text, (int(x1), int(y1) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)
os.makedirs(save_folder, exist_ok=True)
if not filename.endswith((".png", ".jpg", ".jpeg")):
filename += ".png"
filepath = os.path.join(save_folder, filename)
plt.figure(figsize=(12, 8), dpi=300)
plt.imshow(cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB))
plt.title(title, fontsize=16, fontweight="bold")
plt.axis("off")
plt.tight_layout()
plt.savefig(filepath, bbox_inches="tight", dpi=300, format="png", facecolor="white", edgecolor="none")
plt.close()
print(f"High quality image saved: {filepath}")
return filepath