import math 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 car_detection_model = YOLO(r"car_detection.pt") part_detection_model = YOLO(r"car_part.pt") 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, } # Load the models # Define viewing angle rules for scoring (using only the available classes) viewing_angle_rules = { # "Front": { # "must_be_visible": ["Front-bumper", "Grille", "Headlight", "Windshield", "License-plate", "Mirror", "Hood"], # "optional_parts": ["Front-wheel", "Front-window", "Fender", "Quarter-panel", "Rocker-panel"], # "conflict_parts": ["Tail-light", "Back-bumper", "Back-window", "Back-windshield", "Back-wheel", "Trunk"] # }, "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", ], "conflict_parts": ["Grille", "Trunk", "Windshield", "License-plate"], }, "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": { # "must_be_visible": ["Back-bumper", "Tail-light", "Trunk", "Back-windshield", "License-plate", "Roof"], # "optional_parts": ["Back-window", "Rocker-panel", "Mirror", "Back-wheel", "Back-door"], # "conflict_parts": ["Front-bumper", "Headlight", "Grille", "Front-wheel", "Front-door", "Windshield", "Hood", "Fender", "Quarter-panel", "Front-window"] # }, "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", ], "conflict_parts": ["Grille", "Trunk", "Windshield", "License-plate"], }, "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"): """ Compute vehicle side direction using the Mirror as the reference. Group A: ["Windshield", "Hood", "Headlight", "Front-bumper", "Front-wheel"] * If these parts are to the right of the Mirror (positive x-offset), they vote "Right side view". * If to the left, they vote "Left side view". Group B: ["Back-wheel", "Back-door", "Quarter-panel", "Rocker-panel", "Back-window"] * If these parts are to the left of the Mirror (negative x-offset), they vote "Right side view". * If to the right, they vote "Left side view". """ if fixed_label not in detected: # print(f"Reference label '{fixed_label}' not 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)) # print(f"{part} center: {part_center}, x-offset from Mirror: {offset}") 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)) # print(f"{part} center: {part_center}, x-offset from Mirror: {offset}") 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" # NEW LOGIC if voteA and voteB: if voteA == voteB: direction = f"{voteA} side view" reason = f"Both groups agreed on {voteA} side view." else: # Conflict → prioritize group with more datapoints if len(offsets_A) > len(offsets_B): direction = f"{voteA} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group A due to more datapoints." elif len(offsets_B) > len(offsets_A): direction = f"{voteB} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group B due to more datapoints." else: # Equal datapoints → default to Group B (your old rule) direction = f"{voteB} side view" reason = f"Equal datapoints. Falling back to Group B's vote ({voteB})." elif voteA: # Only Group A has datapoints direction = f"{voteA} side view" reason = f"Only Group A datapoints ({len(offsets_A)}). Voting {voteA}." elif voteB: # Only Group B has datapoints direction = f"{voteB} side view" reason = f"Only Group B datapoints ({len(offsets_B)}). Voting {voteB}." else: direction = "Unknown" reason = "No sufficient data." # --- CONSENSUS FLAG LOGIC --- if voteA and voteB: consensus = voteA == voteB # True if same, False if conflict elif voteA or voteB: consensus = True # only one group voted else: consensus = None # no votes at al # print(f"[Mirror] Final guess: {direction}. Reason: {reason}") return direction, radial_lines, consensus def compute_direction_front_wheel_refined(detected, fixed_label="Front-wheel"): """ Compute vehicle side direction using the Front Wheel as the reference. Group A: ["Front-bumper", "Headlight", "Fender", "Grille"] * If these parts are to the left of the Front Wheel (negative x-offset), they vote "Left side view". * If to the right, they vote "Right side view". Group B: ["Front-door", "Back-door", "Rocker-panel", "Front-window", "Back-window", "Back-wheel", "Quarter-panel", "Mirror", "Windshield"] * If these parts are to the right of the Front Wheel (positive x-offset), they vote "Left side view". * If to the left, they vote "Right side view". """ if fixed_label not in detected: # print(f"Reference label '{fixed_label}' not 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)) # print(f"{part} center: {part_center}, x-offset from Front Wheel: {offset}") 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)) # print(f"{part} center: {part_center}, x-offset from Front Wheel: {offset}") 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" # NEW LOGIC if voteA and voteB: if voteA == voteB: direction = f"{voteA} side view" reason = f"Both groups agreed on {voteA} side view." else: # Conflict → prioritize group with more datapoints if len(offsets_A) > len(offsets_B): direction = f"{voteA} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group A due to more datapoints." elif len(offsets_B) > len(offsets_A): direction = f"{voteB} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group B due to more datapoints." else: # Equal datapoints → default to Group B (your old rule) direction = f"{voteB} side view" reason = f"Equal datapoints. Falling back to Group B's vote ({voteB})." elif voteA: # Only Group A has datapoints direction = f"{voteA} side view" reason = f"Only Group A datapoints ({len(offsets_A)}). Voting {voteA}." elif voteB: # Only Group B has datapoints direction = f"{voteB} side view" reason = f"Only Group B datapoints ({len(offsets_B)}). Voting {voteB}." else: direction = "Unknown" reason = "No sufficient data." # --- CONSENSUS FLAG LOGIC --- if voteA and voteB: consensus = voteA == voteB # True if same, False if conflict elif voteA or voteB: consensus = True # only one group voted else: consensus = None # no votes at al # print(f"[Front-wheel] Final guess: {direction}. Reason: {reason}") return direction, radial_lines, consensus def compute_direction_back_wheel_refined(detected, fixed_label="Back-wheel"): """ Compute vehicle side direction using the Back Wheel as the reference. Group A: ["Back-bumper", "Tail-light", "Quarter-panel", "Trunk"] * If these parts are to the left of the Back Wheel (negative x-offset), they vote "Right side view". * If to the right, they vote "Left side view". Group B: ["Front-wheel", "Rocker-panel", "Back-door", "Front-door", "Back-window", "Front-window", "Mirror", "Fender"] * If these parts are to the left of the Back Wheel (negative x-offset), they vote "Left side view". * If to the right, they vote "Right side view". """ if fixed_label not in detected: # print(f"Reference label '{fixed_label}' not 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) # For Group A: parts to the left (negative offset) vote "Right side view", else "Left side view" 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" # NEW LOGIC if voteA and voteB: if voteA == voteB: direction = f"{voteA} side view" reason = f"Both groups agreed on {voteA} side view." else: # Conflict → prioritize group with more datapoints if len(offsets_A) > len(offsets_B): direction = f"{voteA} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group A due to more datapoints." elif len(offsets_B) > len(offsets_A): direction = f"{voteB} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group B due to more datapoints." else: # Equal datapoints → default to Group B (your old rule) direction = f"{voteB} side view" reason = f"Equal datapoints. Falling back to Group B's vote ({voteB})." elif voteA: # Only Group A has datapoints direction = f"{voteA} side view" reason = f"Only Group A datapoints ({len(offsets_A)}). Voting {voteA}." elif voteB: # Only Group B has datapoints direction = f"{voteB} side view" reason = f"Only Group B datapoints ({len(offsets_B)}). Voting {voteB}." else: direction = "Unknown" reason = "No sufficient data." # --- CONSENSUS FLAG LOGIC --- if voteA and voteB: consensus = voteA == voteB # True if same, False if conflict elif voteA or voteB: consensus = True # only one group voted else: consensus = None # no votes at al # print(f"[Back-wheel] Final guess: {direction}. Reason: {reason}") return direction, radial_lines, consensus def compute_direction_headlight_refined(detected, fixed_label="Headlight"): """ Compute vehicle side direction using the Headlight as the reference. 'Fender', 'Windshield', 'Headlight', 'Grille', 'Front-wheel', 'Hood' Group A: ["Front-wheel", "Fender", "Mirror", "Rocker-panel", "Front-door", "Front-window"] * If these parts are to the left of the Headlight (negative x-offset), they vote "Right side view". * If to the right, they vote "Left side view". Group B: ["Front-bumper", "Grille", "Hood", "Windshield"] * If these parts are to the right of the Headlight (positive x-offset), they vote "Right side view". * If to the left, they vote "Left side view". """ if fixed_label not in detected: # print(f"Reference label '{fixed_label}' not 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)) # print(f"{part} center: {part_center}, x-offset from Headlight: {offset}") 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)) # print(f"{part} center: {part_center}, x-offset from Headlight: {offset}") 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" # NEW LOGIC if voteA and voteB: if voteA == voteB: direction = f"{voteA} side view" reason = f"Both groups agreed on {voteA} side view." else: # Conflict → prioritize group with more datapoints if len(offsets_A) > len(offsets_B): direction = f"{voteA} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group A due to more datapoints." elif len(offsets_B) > len(offsets_A): direction = f"{voteB} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group B due to more datapoints." else: # Equal datapoints → default to Group B (your old rule) direction = f"{voteB} side view" reason = f"Equal datapoints. Falling back to Group B's vote ({voteB})." elif voteA: # Only Group A has datapoints direction = f"{voteA} side view" reason = f"Only Group A datapoints ({len(offsets_A)}). Voting {voteA}." elif voteB: # Only Group B has datapoints direction = f"{voteB} side view" reason = f"Only Group B datapoints ({len(offsets_B)}). Voting {voteB}." else: direction = "Unknown" reason = "No sufficient data." # --- CONSENSUS FLAG LOGIC --- if voteA and voteB: consensus = voteA == voteB # True if same, False if conflict elif voteA or voteB: consensus = True # only one group voted else: consensus = None # no votes at al # print(f"[Headlight] Final guess: {direction}. Reason: {reason}") return direction, radial_lines, consensus def compute_direction_tail_refined(detected, fixed_label="Tail-light"): """ Compute vehicle side direction using the Tail-light as the reference. Group A: ["Trunk", "Back-bumper", "Back-windshield"] * If these parts are to the left (negative x-offset) of the Tail-light, vote "Right side view"; if to the right, vote "Left side view". Group B: ["Back-wheel", "Quarter-panel", "Back-door", "Back-window"] * If these parts are to the right (positive x-offset) of the Tail-light, vote "Right side view"; if to the left, vote "Left side view". """ if fixed_label not in detected: # print(f"Reference label '{fixed_label}' not 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)) # print(f"{part} center: {part_center}, x-offset from Tail-light: {offset}") 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)) # print(f"{part} center: {part_center}, x-offset from Tail-light: {offset}") voteA = None voteB = None if offsets_A: avg_A = sum(offsets_A) / len(offsets_A) # print(f"Average Group A offset (Tail-light): {avg_A}") voteA = "Right" if avg_A < 0 else "Left" if offsets_B: avg_B = sum(offsets_B) / len(offsets_B) # print(f"Average Group B offset (Tail-light): {avg_B}") voteB = "Right" if avg_B > 0 else "Left" # NEW LOGIC if voteA and voteB: if voteA == voteB: direction = f"{voteA} side view" reason = f"Both groups agreed on {voteA} side view." else: # Conflict → prioritize group with more datapoints if len(offsets_A) > len(offsets_B): direction = f"{voteA} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group A due to more datapoints." elif len(offsets_B) > len(offsets_A): direction = f"{voteB} side view" reason = f"Conflict: Group A ({len(offsets_A)}) vs Group B ({len(offsets_B)}). Prioritizing Group B due to more datapoints." else: # Equal datapoints → default to Group B (your old rule) direction = f"{voteB} side view" reason = f"Equal datapoints. Falling back to Group B's vote ({voteB})." elif voteA: # Only Group A has datapoints direction = f"{voteA} side view" reason = f"Only Group A datapoints ({len(offsets_A)}). Voting {voteA}." elif voteB: # Only Group B has datapoints direction = f"{voteB} side view" reason = f"Only Group B datapoints ({len(offsets_B)}). Voting {voteB}." else: direction = "Unknown" reason = "No sufficient data." # --- CONSENSUS FLAG LOGIC --- if voteA and voteB: consensus = voteA == voteB # True if same, False if conflict elif voteA or voteB: consensus = True # only one group voted else: consensus = None # no votes at al # print(f"[Tail-light] Final guess: {direction}. Reason: {reason}") return direction, radial_lines, consensus def determine_vehicle_directions(detected): directions = {} radial_lines_all = {} consensus_flags = {} # Mirror-based direction 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-light-based direction 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-based direction 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-based direction 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-based direction 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): """ Returns: directions: {"Selected": ""} or {"Selected":"Not Applicable"/"Unknown"} detected_parts: set(...) need_review: bool top_2_predictions: [top1_key_lower, top2_key_lower_or_False] score_map: dict mapping angle_key_lower -> stage2_score (for all angles) """ need_review = False image_array = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR) car_results = car_detection_model(image_array) # --- CAR DETECTION VISUALIZATION --- car_detections_to_plot = [] main_car_bbox = None best_area = 0 num = 0 for result in car_results: for box in result.boxes.data: num += 1 conf = float(box[4]) if conf < 0.50: # confidence threshold for cars 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) # print("Number of car boxes:", num) # plot_detections(image_array, car_detections_to_plot, title="Car Detections") if main_car_bbox is None: # keep shape consistent return pil_image, {}, {} # --- PART DETECTION VISUALIZATION --- 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] ) ignored_parts = [] kept_parts = [] for result in results: for box in result.boxes.data: conf = float(box[4]) if conf < 0.65: ignored_parts.append((result.names[int(box[5])], "low_conf", conf)) 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 # lookup per-part min ratio min_ratio = MIN_RATIO.get(label, MIN_RATIO["default"]) # filter: too small relative to car if part_ratio < min_ratio and conf < 0.8: ignored_parts.append( (label, "too_small", float(part_ratio), float(conf)) ) continue # filter: truncated parts touching bbox edges 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: # stricter rule for edge-cut parts ignored_parts.append( (label, "truncated_edge", float(part_ratio), float(conf)) ) continue # keep part detected.setdefault(label, []).append((x1, y1, x2, y2)) part_detections_to_plot.append((x1, y1, x2, y2, label, conf)) kept_parts.append((label, float(part_ratio), float(conf))) # plot_detections(image_array, part_detections_to_plot, title="Part Detections") 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): """ Returns: directions: {"Selected": ""} or {"Selected":"Not Applicable"/"Unknown"} detected_parts: set(...) need_review: bool top_2_predictions: [top1_key_lower, top2_key_lower_or_False] score_map: dict mapping angle_key_lower -> stage2_score (for all angles) """ need_review = False # plot_detections(image_array, part_detections_to_plot, title="Part Detections") detected_parts = set(detected.keys()) # print(f"[Summary] Detected Parts (inside main car): {detected_parts}") # Compute scores for each angle angle_scores = [] critical_parts = { # "Front": {"Front-bumper", "Headlight", "Windshield", "Grille"}, "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": {"Tail-light", "Trunk", "Back-windshield", "Back-bumper"}, "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 # print(f"[Summary] Angle: {angle}, Stage2 Score: {stage2_score:.2f}") angle_scores.append((angle, critical_ratio, stage2_score)) # Build score_map for all angles (lowercase keys) score_map = {angle.lower(): score for (angle, _, score) in angle_scores} # ---- Stage-1 selection: ALWAYS pick top1; pick top2 only if strictly lower and >= threshold ---- # eps = 1e-6 # view_diff_thresold = 0.20 # critical_thresold = 0.75 stage_2_thresold = 0.80 # print(angle_scores) sorted_all = sorted(angle_scores, key=lambda x: x[2], reverse=True) top_2_predictions = ["", ""] # [top1_key_lower, top2_key_lower_or_False] # print("sorted",sorted_all) if sorted_all: top1 = sorted_all[0] top1_key = top1[0].lower() top1_score = top1[2] top1_cric_score = top1[1] top_2_predictions[0] = top1_key # find second candidate: strictly lower than top1 and >= threshold second_key = "" for angle, crit, score in sorted_all[1:]: # print(angle,score) if score < top1_score and score >= stage_2_thresold: # print(angle) second_key = angle.lower() break top_2_predictions[1] = second_key # print(f"[Summary] Stage1 Top1: {top_2_predictions[0]}, Stage1 Top2 (or False): {top_2_predictions[1]}") # Existing selection logic to pick the best_angle (unchanged) 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, best_critical, best_stage2 = 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, best_critical, best_stage2 = max(candidates, key=lambda x: x[1]) else: best_angle, best_critical, best_stage2 = max( angle_scores, key=lambda x: x[2] ) # print(f"[Summary] Final Viewing Angle (from scoring): {best_angle}") # Stage-2 direction (geometric) if best_angle in ["Front", "Rear"]: directions = {"Selected": best_angle} else: directions_all, radial_lines_all, consensus_all = determine_vehicle_directions( detected ) # --- NEW CONSENSUS PRIORITIZATION --- # 1. Check if any anchor has consensus=True 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: # Pick the first consensus-true vote (or implement a tie-breaker if needed) chosen_ref, chosen_dir = next(iter(consensus_votes.items())) # print(f"[Consensus Override] Using {chosen_ref} vote because consensus=True → {chosen_dir}") majority_side = chosen_dir.split()[0] else: # --- FALL BACK TO ORIGINAL VOTING LOGIC --- 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" # print(f"[Stage-2 Majority Side] {majority_side}") 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} # print(f"[Summary] Final side classification: {final_side_classification}.") return directions, detected_parts, need_review, top_2_predictions, score_map def deskew_image(pil_image: Image.Image) -> Image.Image: """ Correct skew/rotation of an input PIL Image using Hough Line Transform. Args: pil_image (PIL.Image.Image): Input image. Returns: PIL.Image.Image: Deskewed image. """ # Convert PIL to OpenCV (BGR) img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR) # Convert to grayscale gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # Detect edges edges = cv2.Canny(gray, 50, 150, apertureSize=3) # Hough Line Transform lines = cv2.HoughLines(edges, 1, np.pi / 180, 200) angles = [] if lines is not None: for rho, theta in lines[:, 0]: angle = (theta * 180 / np.pi) - 90 # Normalize angle to [-90, 90] if angle < -90: angle += 180 if angle > 90: angle -= 180 angles.append(angle) # Use median angle median_angle = np.median(angles) if len(angles) > 0 else 0 # print("Estimated angle:", median_angle) # Rotate image to deskew (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 ) # Convert back to PIL (RGB) return Image.fromarray(cv2.cvtColor(rotated, cv2.COLOR_BGR2RGB)) async def yolo_rule_based_classification(pil_image, image_name, img_file: UploadFile): """ Returns: mapped_primary: canonical label (Stage-1 top1 mapped, with Stage-2 side enforced) final_review: bool final_secondaries: two-slot list [mapped_top1, mapped_top2_or_False] - mapped_top2 is included only if Stage-1 top2 exists AND its Stage-1 score >= 0.85 """ 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 = {} for rotation in rotations: rotated_image = ( pil_image.rotate(rotation, expand=True) if rotation != 0 else pil_image ) pil_image, detected, detections = find_best_combination(rotated_image) 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 View", "right": "Passenger Side View", "left side view": "Driver Side View", "right side view": "Passenger Side View", "front right": "Front Passenger Side Corner View", "front left": "Front Driver Side Corner View", "rear right": "Rear Passenger Side Corner View", "rear left": "Rear Driver Side Corner View", "unknown": "NA", } # derive desired_side from Stage-2 raw direction desired_side, opposite_side = None, None if isinstance(best_raw_direction, str): raw = best_raw_direction.lower() if "left" in raw: desired_side = "Driver Side View" opposite_side = "Passenger Side View" elif "right" in raw: desired_side = "Passenger Side View" opposite_side = "Driver Side View" # Stage-1 keys (strings or False) 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: prefer Stage-1 top1 mapped -> enforce Stage-2 side if needed, fallback to Stage-2 raw 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") # corner swap mapping (for side enforcement) corner_swap = { "Front Passenger Side Corner View": "Front Driver Side Corner View", "Front Driver Side Corner View": "Front Passenger Side Corner View", "Rear Passenger Side Corner View": "Rear Driver Side Corner View", "Rear Driver Side Corner View": "Rear Passenger Side Corner View", } # enforce desired_side on mapped_primary if necessary if desired_side and mapped_primary != "NA": if mapped_primary in ("Driver Side View", "Passenger Side View"): if mapped_primary != desired_side: mapped_primary = desired_side elif mapped_primary in corner_swap: if desired_side == "Driver Side View" and "Passenger" in mapped_primary: mapped_primary = corner_swap[mapped_primary] elif desired_side == "Passenger Side View" and "Driver" in mapped_primary: mapped_primary = corner_swap[mapped_primary] # Build final_secondaries strictly from Stage-1 keys: final_secondaries = [False, False] threshold = 0.8 eps = 1e-6 # mapped_top1 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 # enforce side on mapped_top1 if desired_side and mapped_top1 and mapped_top1 in corner_swap: if desired_side == "Driver Side View" and "Passenger" in mapped_top1: mapped_top1 = corner_swap[mapped_top1] elif desired_side == "Passenger Side View" and "Driver" in mapped_top1: mapped_top1 = corner_swap[mapped_top1] elif desired_side and mapped_top1 in ("Driver Side View", "Passenger Side View"): if mapped_top1 != desired_side: mapped_top1 = desired_side final_secondaries[0] = mapped_top1 if mapped_top1 != "NA" else False # mapped_top2 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 - eps): mapped_top2 = viewing_angle_map.get(stage1_top2_key.lower(), "NA") # enforce desired_side if desired_side and mapped_top2 in corner_swap: if desired_side == "Driver Side View" and "Passenger" in mapped_top2: mapped_top2 = corner_swap[mapped_top2] elif desired_side == "Passenger Side View" and "Driver" in mapped_top2: mapped_top2 = corner_swap[mapped_top2] elif desired_side and mapped_top2 in ( "Driver Side View", "Passenger Side View", ): if mapped_top2 != desired_side: mapped_top2 = desired_side else: mapped_top2 = False final_secondaries[1] = mapped_top2 if mapped_top2 and mapped_top2 != "NA" else False # Fallback if primary is NA if mapped_primary == "NA" and final_secondaries[0]: mapped_primary = final_secondaries[0] # Final-review heuristics if fail_count >= 3: final_review = True # collect scores 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)] 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(f"{output_folder}", exist_ok=True) output_file_car = f"car_detection{extension}" output_file_part = f"part_detection{extension}" print("\n") print("output_file_car = ", output_file_car) print("output_file_car = ", output_file_part) print("\n") 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() # Draw bounding boxes and labels with better styling for x1, y1, x2, y2, label, conf in detections: # Thicker green rectangle cv2.rectangle(vis_img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 3) # Better text with background for readability 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, ) # Create save folder os.makedirs(save_folder, exist_ok=True) # Add .png extension if not provided if not filename.endswith((".png", ".jpg", ".jpeg")): filename += ".png" filepath = os.path.join(save_folder, filename) # High quality save with matplotlib 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 def plot_detections(image_array, detections, title="Detections"): """ image_array: numpy BGR image detections: list of tuples (x1, y1, x2, y2, label, conf) """ 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), 2) cv2.putText( vis_img, f"{label} {conf:.2f}", (int(x1), int(y1) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2, ) plt.figure(figsize=(8, 6)) plt.imshow(cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)) plt.title(title) plt.axis("off") plt.show()