Spaces:
Sleeping
Sleeping
| 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": "<angle>"} 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": "<angle>"} 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() | |