Spaces:
Sleeping
Sleeping
| 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 | |