Update app.py

app.py

@@ -25,63 +25,135 @@ class_names = {0: 'With Helmet', 1: 'Without Helmet', 2: 'License Plate'}
 def crop_license_plates(image, detections):
     cropped_plates = []
     
-    if isinstance(image, str):
-        image = Image.open(image)
-    elif isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
+    try:
+        if isinstance(image, str):
+            if not os.path.exists(image):
+                print(f"Error: Image file not found: {image}")
+                return cropped_plates
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        elif not isinstance(image, Image.Image):
+            print(f"Error: Unsupported image type: {type(image)}")
+            return cropped_plates
+        
+        if image.size[0] == 0 or image.size[1] == 0:
+            print("Error: Image has zero dimensions")
+            return cropped_plates
+            
+    except Exception as e:
+        print(f"Error loading image: {e}")
+        return cropped_plates
     
-    for detection in detections:
-        if detection['Object'] == 'License Plate':
-            pos = detection['Position'].strip('()')
-            x1, y1 = map(int, pos.split(', '))
+    for i, detection in enumerate(detections):
+        try:
+            if detection['Object'] != 'License Plate':
+                continue
+                
+            pos_str = detection['Position'].strip('()')
+            if ',' not in pos_str:
+                print(f"Error: Invalid position format for detection {i}: {detection['Position']}")
+                continue
+                
+            x1, y1 = map(int, pos_str.split(', '))
+            
+            dims_str = detection['Dimensions']
+            if 'x' not in dims_str:
+                print(f"Error: Invalid dimensions format for detection {i}: {detection['Dimensions']}")
+                continue
+                
+            width, height = map(int, dims_str.split('x'))
+            
+            if width <= 0 or height <= 0:
+                print(f"Error: Invalid dimensions for detection {i}: {width}x{height}")
+                continue
             
-            dims = detection['Dimensions']
-            width, height = map(int, dims.split('x'))
             x2, y2 = x1 + width, y1 + height
             
-            padding = 10
-            x1 = max(0, x1 - padding)
-            y1 = max(0, y1 - padding)
-            x2 = min(image.width, x2 + padding)
-            y2 = min(image.height, y2 + padding)
+            if x1 < 0 or y1 < 0 or x2 > image.width or y2 > image.height:
+                print(f"Warning: Bounding box extends beyond image boundaries for detection {i}")
+                x1 = max(0, x1)
+                y1 = max(0, y1)
+                x2 = min(image.width, x2)
+                y2 = min(image.height, y2)
+            
+            if x2 <= x1 or y2 <= y1:
+                print(f"Error: Invalid crop coordinates for detection {i}: ({x1},{y1}) to ({x2},{y2})")
+                continue
             
             cropped_plate = image.crop((x1, y1, x2, y2))
+            
+            if cropped_plate.size[0] == 0 or cropped_plate.size[1] == 0:
+                print(f"Error: Cropped image has zero dimensions for detection {i}")
+                continue
+            
             cropped_plates.append({
                 'image': cropped_plate,
                 'confidence': detection['Confidence'],
-                'position': detection['Position']
+                'position': detection['Position'],
+                'crop_coords': f"({x1},{y1}) to ({x2},{y2})"
             })
+            
+        except ValueError as e:
+            print(f"Error parsing coordinates for detection {i}: {e}")
+            continue
+        except Exception as e:
+            print(f"Error cropping license plate {i}: {e}")
+            continue
     
     return cropped_plates
 
 def create_download_files(annotated_image, cropped_plates, detections):
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    
-    os.makedirs("temp", exist_ok=True)
-    
-    annotated_path = f"temp/annotated_image_{timestamp}.jpg"
-    annotated_image.save(annotated_path)
-    
-    plate_paths = []
-    for i, plate_data in enumerate(cropped_plates):
-        plate_path = f"temp/license_plate_{i+1}_{timestamp}.jpg"
-        plate_data['image'].save(plate_path)
-        plate_paths.append(plate_path)
-    
-    report_path = f"temp/detection_report_{timestamp}.csv"
-    if detections:
-        df = pd.DataFrame(detections)
-        df.to_csv(report_path, index=False)
-    
-    zip_path = f"temp/detection_results_{timestamp}.zip"
-    with zipfile.ZipFile(zip_path, 'w') as zipf:
-        zipf.write(annotated_path, f"annotated_image_{timestamp}.jpg")
-        for plate_path in plate_paths:
-            zipf.write(plate_path, os.path.basename(plate_path))
-        if os.path.exists(report_path):
-            zipf.write(report_path, f"detection_report_{timestamp}.csv")
-    
-    return zip_path, annotated_path, plate_paths
+    try:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        
+        os.makedirs("temp", exist_ok=True)
+        
+        annotated_path = f"temp/annotated_image_{timestamp}.jpg"
+        try:
+            annotated_image.save(annotated_path, quality=95)
+        except Exception as e:
+            print(f"Error saving annotated image: {e}")
+            return None, None, []
+        
+        plate_paths = []
+        for i, plate_data in enumerate(cropped_plates):
+            try:
+                plate_path = f"temp/license_plate_{i+1}_{timestamp}.jpg"
+                plate_data['image'].save(plate_path, quality=95)
+                plate_paths.append(plate_path)
+            except Exception as e:
+                print(f"Error saving license plate {i+1}: {e}")
+                continue
+        
+        report_path = f"temp/detection_report_{timestamp}.csv"
+        if detections:
+            try:
+                df = pd.DataFrame(detections)
+                df.to_csv(report_path, index=False)
+            except Exception as e:
+                print(f"Error creating detection report: {e}")
+                report_path = None
+        
+        zip_path = f"temp/detection_results_{timestamp}.zip"
+        try:
+            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+                if os.path.exists(annotated_path):
+                    zipf.write(annotated_path, f"annotated_image_{timestamp}.jpg")
+                for plate_path in plate_paths:
+                    if os.path.exists(plate_path):
+                        zipf.write(plate_path, os.path.basename(plate_path))
+                if report_path and os.path.exists(report_path):
+                    zipf.write(report_path, f"detection_report_{timestamp}.csv")
+        except Exception as e:
+            print(f"Error creating ZIP file: {e}")
+            return None, annotated_path, plate_paths
+        
+        return zip_path, annotated_path, plate_paths
+        
+    except Exception as e:
+        print(f"Error in create_download_files: {e}")
+        return None, None, []
 
 def yoloV8_func(
     image=None, 
@@ -137,15 +209,19 @@ def yoloV8_func(
     download_files = None
     
     if crop_plates and detections:
-        cropped_plates = crop_license_plates(image, detections)
-        license_plate_gallery = [plate_data['image'] for plate_data in cropped_plates]
-        
-        if cropped_plates or detections:
-            try:
+        try:
+            cropped_plates = crop_license_plates(image, detections)
+            license_plate_gallery = [plate_data['image'] for plate_data in cropped_plates]
+            
+            if cropped_plates or detections:
                 download_files, _, _ = create_download_files(annotated_image, cropped_plates, detections)
-            except Exception as e:
-                print(f"Error creating download files: {e}")
-                download_files = None
+                if download_files is None:
+                    print("Warning: Could not create download files")
+        except Exception as e:
+            print(f"Error in license plate processing: {e}")
+            cropped_plates = []
+            license_plate_gallery = []
+            download_files = None
     
     # Create stats text
     stats_text = ""